Add tests for import hooks in the driver wrappers

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index bb08bf6..7533ebf 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -42,9 +42,13 @@
     /* Count the amount of times one of the key management driver functions
      * is called. */
     unsigned long hits;
+    /* Record the source of the function call. */
+    psa_key_location_t source;
 } mbedtls_test_driver_key_management_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0 }
+/* 0x800000 is a vendor-specific location, unused by the PSA, overwritten
+ * in tests that expect a different value. */
+#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0x800000 }
 static inline mbedtls_test_driver_key_management_hooks_t
     mbedtls_test_driver_key_management_hooks_init( void )
 {