Add key export on successful key agreement

More sanity checks on key coming out of key agreement.

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 62fa455..7b082a3 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -693,7 +693,13 @@
     size_t public_key_length;
     uint8_t output[1024];
     size_t output_length;
+
+    uint8_t *exported = NULL;
+    size_t exported_size = 0;
+    size_t exported_length = 0;
+
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_attributes_t export_attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     mbedtls_svc_key_id_t shared_secret_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t shared_secret_attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -746,6 +752,28 @@
     if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
         /* The key has been destroyed. */
         status = PSA_SUCCESS;
+        goto exit;
+    } else if (status == PSA_SUCCESS) {
+
+        status = psa_get_key_attributes(shared_secret_id, &export_attributes);
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            status = PSA_SUCCESS;
+            goto exit;
+        }
+
+        exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(psa_get_key_type(&export_attributes),
+                                                   psa_get_key_bits(&export_attributes));
+        TEST_CALLOC(exported, exported_size);
+
+        status = psa_export_key(shared_secret_id, exported, exported_size, &exported_length);
+
+        if (key_destroyable && status == PSA_ERROR_INVALID_HANDLE) {
+            /* The key has been destroyed. */
+            status = PSA_SUCCESS;
+        }
+
+        PSA_ASSERT(status);
     }
 
 exit:
@@ -754,12 +782,14 @@
      * thus reset them as required.
      */
     psa_reset_key_attributes(&attributes);
+    psa_reset_key_attributes(&export_attributes);
 
     /* Make sure to reset and free derived key attributes and slot. */
     psa_reset_key_attributes(&shared_secret_attributes);
     psa_destroy_key(shared_secret_id);
 
     mbedtls_free(public_key);
+    mbedtls_free(exported);
     return status;
 }