Restructure the crypto secure element tests, re-add previous tests

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 323ad2d..9dd7f88 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -136,6 +136,7 @@
 add_test_suite(psa_crypto_metadata)
 add_test_suite(psa_crypto_persistent_key)
 add_test_suite(psa_crypto_se_driver_hal)
+add_test_suite(psa_crypto_se_driver_hal_mocks)
 add_test_suite(psa_crypto_slot_management)
 add_test_suite(psa_its)
 add_test_suite(shax)
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.data b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
index ef40e7e..035d79d 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -27,23 +27,11 @@
 Register SE driver: maximum number of drivers
 register_max:
 
-Key importing mock test
-mock_import:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+Key creation smoke test (no p_allocate)
+key_creation_import_export:-1
 
-Key importing mock test: alloc failed
-mock_import:PSA_ERROR_HARDWARE_FAILURE:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE
+Key creation smoke test (p_allocate allows all slots)
+key_creation_import_export:0
 
-Key importing mock test: import failed
-mock_import:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
-
-Key exporting mock test
-mock_export:PSA_SUCCESS:PSA_SUCCESS
-
-Key exporting mock test: export failed
-mock_export:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
-
-Key generating mock test
-mock_generate:PSA_SUCCESS:PSA_SUCCESS
-
-Key generating mock test: generating failed
-mock_generate:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
+Key creation smoke test (p_allocate allows 1 slot)
+key_creation_import_export:ARRAY_LENGTH( ram_slots ) - 1
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index bfadf3d..33a8c4e 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -11,141 +11,77 @@
  * This is probably a bug in the library. */
 #define PSA_ERROR_DETECTED_BY_DRIVER ((psa_status_t)( -500 ))
 
-static struct
+#define RAM_MAX_KEY_SIZE 64
+typedef struct
 {
-    uint16_t called;
-    psa_key_slot_number_t key_slot;
-    psa_key_type_t type;
-    psa_key_usage_t usage;
-    size_t bits;
-    size_t extra_size;
-    size_t pubkey_out_size;
-    psa_status_t return_value;
-} mock_generate_data;
-
-static struct
-{
-    uint16_t called;
-    psa_key_slot_number_t slot_number;
     psa_key_lifetime_t lifetime;
     psa_key_type_t type;
-    psa_algorithm_t algorithm;
-    psa_key_usage_t usage;
-    size_t data_length;
-    psa_status_t return_value;
-} mock_import_data;
+    size_t bits;
+    uint8_t content[RAM_MAX_KEY_SIZE];
+} ram_slot_t;
+static ram_slot_t ram_slots[16];
 
-static struct
-{
-    uint16_t called;
-    psa_key_slot_number_t slot_number;
-    size_t data_size;
-    psa_status_t return_value;
-} mock_export_data;
+static uint8_t ram_min_slot = 0;
 
-static struct
+static void ram_slots_reset( void )
 {
-    uint16_t called;
-    psa_status_t return_value;
-} mock_allocate_data;
-
-static struct
-{
-    uint16_t called;
-    psa_key_slot_number_t slot_number;
-    psa_status_t return_value;
-} mock_destroy_data;
-
-static void mock_teardown( void )
-{
-    memset( &mock_import_data, 0, sizeof( mock_import_data ) );
-    memset( &mock_export_data, 0, sizeof( mock_export_data ) );
-    memset( &mock_allocate_data, 0, sizeof( mock_allocate_data ) );
-    memset( &mock_destroy_data, 0, sizeof( mock_destroy_data ) );
-    memset( &mock_generate_data, 0, sizeof( mock_generate_data ) );
+    memset( ram_slots, 0, sizeof( ram_slots ) );
+    ram_min_slot = 0;
 }
 
-static psa_status_t mock_generate( psa_key_slot_number_t key_slot,
-                                   psa_key_type_t type,
-                                   psa_key_usage_t usage,
-                                   size_t bits,
-                                   const void *extra,
-                                   size_t extra_size,
-                                   uint8_t *p_pubkey_out,
-                                   size_t pubkey_out_size,
-                                   size_t *p_pubkey_length )
+static psa_status_t ram_import( psa_key_slot_number_t slot_number,
+                                psa_key_lifetime_t lifetime,
+                                psa_key_type_t type,
+                                psa_algorithm_t algorithm,
+                                psa_key_usage_t usage,
+                                const uint8_t *p_data,
+                                size_t data_length )
 {
-    (void) extra;
-    (void) p_pubkey_out;
-    (void) p_pubkey_length;
-
-    mock_generate_data.called++;
-    mock_generate_data.key_slot = key_slot;
-    mock_generate_data.type = type;
-    mock_generate_data.usage = usage;
-    mock_generate_data.bits = bits;
-    mock_generate_data.extra_size = extra_size;
-    mock_generate_data.pubkey_out_size = pubkey_out_size;
-
-    return( mock_generate_data.return_value );
+    if( slot_number > ARRAY_LENGTH( ram_slots ) )
+        return( PSA_ERROR_DETECTED_BY_DRIVER );
+    if( data_length > sizeof( ram_slots[slot_number].content ) )
+        return( PSA_ERROR_INSUFFICIENT_STORAGE );
+    ram_slots[slot_number].lifetime = lifetime;
+    ram_slots[slot_number].type = type;
+    ram_slots[slot_number].bits = PSA_BYTES_TO_BITS( data_length );
+    (void) algorithm;
+    (void) usage;
+    memcpy( ram_slots[slot_number].content, p_data, data_length );
+    return( PSA_SUCCESS );
 }
 
-static psa_status_t mock_import( psa_key_slot_number_t slot_number,
-                                 psa_key_lifetime_t lifetime,
-                                 psa_key_type_t type,
-                                 psa_algorithm_t algorithm,
-                                 psa_key_usage_t usage,
-                                 const uint8_t *p_data,
-                                 size_t data_length )
+psa_status_t ram_export( psa_key_slot_number_t slot_number,
+                         uint8_t *p_data,
+                         size_t data_size,
+                         size_t *p_data_length )
 {
-    (void) p_data;
-
-    mock_import_data.called++;
-    mock_import_data.slot_number = slot_number;
-    mock_import_data.lifetime = lifetime;
-    mock_import_data.type = type;
-    mock_import_data.algorithm = algorithm;
-    mock_import_data.usage = usage;
-    mock_import_data.data_length = data_length;
-
-    return( mock_import_data.return_value );
+    size_t actual_size;
+    if( slot_number > ARRAY_LENGTH( ram_slots ) )
+        return( PSA_ERROR_DETECTED_BY_DRIVER );
+    actual_size = PSA_BITS_TO_BYTES( ram_slots[slot_number].bits );
+    if( actual_size > data_size )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+    *p_data_length = actual_size;
+    memcpy( p_data, ram_slots[slot_number].content, actual_size );
+    return( PSA_SUCCESS );
 }
 
-psa_status_t mock_export( psa_key_slot_number_t slot_number,
-                          uint8_t *p_data,
-                          size_t data_size,
-                          size_t *p_data_length )
+psa_status_t ram_destroy( psa_key_slot_number_t slot_number )
 {
-    (void) p_data;
-    (void) p_data_length;
-
-    mock_export_data.called++;
-    mock_export_data.slot_number = slot_number;
-    mock_export_data.data_size = data_size;
-
-    return( mock_export_data.return_value );
+    if( slot_number > ARRAY_LENGTH( ram_slots ) )
+        return( PSA_ERROR_DETECTED_BY_DRIVER );
+    memset( &ram_slots[slot_number], 0, sizeof( ram_slots[slot_number] ) );
+    return( PSA_SUCCESS );
 }
 
-psa_status_t mock_allocate( const psa_key_attributes_t *attributes,
+psa_status_t ram_allocate( const psa_key_attributes_t *attributes,
                            const psa_drv_se_slot_usage_t *slot_usage,
                            psa_key_slot_number_t *slot_number )
 {
     (void) attributes;
-    (void) slot_usage;
-    (void) slot_number;
-
-    mock_allocate_data.called++;
-    *slot_number = 0;
-
-    return( mock_allocate_data.return_value );
-}
-
-psa_status_t mock_destroy( psa_key_slot_number_t slot_number )
-{
-    mock_destroy_data.called++;
-    mock_destroy_data.slot_number = slot_number;
-
-    return( mock_destroy_data.return_value );
+    return( psa_drv_cb_find_free_slot( slot_usage,
+                                       ram_min_slot, -1,
+                                       slot_number ) );
 }
 
 /* END_HEADER */
@@ -221,59 +157,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void mock_import( int mock_alloc_return_value,
-                  int mock_import_return_value,
-                  int expected_result)
-{
-    psa_drv_se_t driver;
-    psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
-    psa_key_id_t id = 1;
-    psa_key_handle_t handle = 0;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
-
-    mock_allocate_data.return_value = mock_alloc_return_value;
-    mock_import_data.return_value = mock_import_return_value;
-    memset( &driver, 0, sizeof( driver ) );
-    memset( &key_management, 0, sizeof( key_management ) );
-    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
-    driver.key_management = &key_management;
-    key_management.slot_count = 1;
-    key_management.p_import = mock_import;
-    key_management.p_destroy = mock_destroy;
-    key_management.p_allocate = mock_allocate;
-
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
-    PSA_ASSERT( psa_crypto_init( ) );
-
-    psa_set_key_id( &attributes, id );
-    psa_set_key_lifetime( &attributes, lifetime );
-    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
-    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
-    TEST_ASSERT( psa_import_key( &attributes,
-                                key_material, sizeof( key_material ),
-                                &handle ) == expected_result );
-
-    TEST_ASSERT( mock_allocate_data.called == 1 );
-    TEST_ASSERT( mock_import_data.called ==
-                  (expected_result == PSA_SUCCESS? 1 : 0) );
-    TEST_ASSERT( mock_import_data.type ==
-                  (expected_result == PSA_SUCCESS? PSA_KEY_TYPE_RAW_DATA : 0));
-
-    if( expected_result == PSA_SUCCESS )
-    {
-        PSA_ASSERT( psa_destroy_key( handle ) );
-        TEST_ASSERT( mock_destroy_data.called == 1 );
-    }
-exit:
-    PSA_DONE( );
-    mock_teardown( );
-}
-/* END_CASE */
-
-/* BEGIN_CASE */
-void mock_export( int mock_export_return_value, int expected_result )
+void key_creation_import_export( int min_slot )
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
@@ -284,17 +168,21 @@
     const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
     uint8_t exported[sizeof( key_material )];
     size_t exported_length;
+    psa_key_slot_number_t expected_slot = ( min_slot >= 0 ? min_slot : 0 );
 
-    mock_export_data.return_value = mock_export_return_value;
     memset( &driver, 0, sizeof( driver ) );
     memset( &key_management, 0, sizeof( key_management ) );
     driver.hal_version = PSA_DRV_SE_HAL_VERSION;
     driver.key_management = &key_management;
-    key_management.slot_count = 1;
-    key_management.p_import = mock_import;
-    key_management.p_export = mock_export;
-    key_management.p_destroy = mock_destroy;
-    key_management.p_allocate = mock_allocate;
+    key_management.slot_count = ARRAY_LENGTH( ram_slots );
+    key_management.p_import = ram_import;
+    key_management.p_export = ram_export;
+    key_management.p_destroy = ram_destroy;
+    if( min_slot >= 0 )
+    {
+        key_management.p_allocate = ram_allocate;
+        ram_min_slot = min_slot;
+    }
 
     PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
@@ -307,54 +195,22 @@
                                 key_material, sizeof( key_material ),
                                 &handle ) );
 
-    TEST_ASSERT( psa_export_key( handle,
-                                exported, sizeof( exported ),
-                                &exported_length ) == expected_result );
+    /* Test that the key was created in the designated slot. */
+    TEST_ASSERT( ram_slots[expected_slot].type == PSA_KEY_TYPE_RAW_DATA );
 
-    TEST_ASSERT( mock_export_data.called == 1 );
+    PSA_ASSERT( psa_export_key( handle,
+                                exported, sizeof( exported ),
+                                &exported_length ) );
+    ASSERT_COMPARE( key_material, sizeof( key_material ),
+                    exported, exported_length );
 
     PSA_ASSERT( psa_destroy_key( handle ) );
 
-    TEST_ASSERT( mock_destroy_data.called == 1 );
+    /* Test that the key has been erased from the designated slot. */
+    TEST_ASSERT( ram_slots[expected_slot].type == 0 );
 
 exit:
     PSA_DONE( );
-    mock_teardown( );
-}
-/* END_CASE */
-
-/* BEGIN_CASE */
-void mock_generate( int mock_generate_return_value, int expected_result )
-{
-    psa_drv_se_t driver;
-    psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = 2;
-    psa_key_id_t id = 1;
-    psa_key_handle_t handle = 0;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-
-    mock_generate_data.return_value = mock_generate_return_value;
-    memset( &driver, 0, sizeof( driver ) );
-    memset( &key_management, 0, sizeof( key_management ) );
-    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
-    driver.key_management = &key_management;
-    key_management.slot_count = 1;
-    key_management.p_generate = mock_generate;
-    key_management.p_destroy = mock_destroy;
-    key_management.p_allocate = mock_allocate;
-
-    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
-    PSA_ASSERT( psa_crypto_init( ) );
-
-    psa_set_key_id( &attributes, id );
-    psa_set_key_lifetime( &attributes, lifetime );
-    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
-    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
-    TEST_ASSERT( psa_generate_key( &attributes, &handle ) == expected_result );
-    TEST_ASSERT( mock_generate_data.called == 1 );
-
-exit:
-    PSA_DONE( );
-    mock_teardown( );
+    ram_slots_reset( );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
new file mode 100644
index 0000000..7cdbc71
--- /dev/null
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
@@ -0,0 +1,20 @@
+Key importing mock test
+mock_import:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+
+Key importing mock test: alloc failed
+mock_import:PSA_ERROR_HARDWARE_FAILURE:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE
+
+Key importing mock test: import failed
+mock_import:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
+
+Key exporting mock test
+mock_export:PSA_SUCCESS:PSA_SUCCESS
+
+Key exporting mock test: export failed
+mock_export:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
+
+Key generating mock test
+mock_generate:PSA_SUCCESS:PSA_SUCCESS
+
+Key generating mock test: generating failed
+mock_generate:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
new file mode 100644
index 0000000..58f8aa4
--- /dev/null
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
@@ -0,0 +1,288 @@
+/* BEGIN_HEADER */
+#include "psa_crypto_helpers.h"
+#include "psa/crypto_se_driver.h"
+
+#include "psa_crypto_se.h"
+
+static struct
+{
+    uint16_t called;
+    psa_key_slot_number_t key_slot;
+    psa_key_type_t type;
+    psa_key_usage_t usage;
+    size_t bits;
+    size_t extra_size;
+    size_t pubkey_out_size;
+    psa_status_t return_value;
+} mock_generate_data;
+
+static struct
+{
+    uint16_t called;
+    psa_key_slot_number_t slot_number;
+    psa_key_lifetime_t lifetime;
+    psa_key_type_t type;
+    psa_algorithm_t algorithm;
+    psa_key_usage_t usage;
+    size_t data_length;
+    psa_status_t return_value;
+} mock_import_data;
+
+static struct
+{
+    uint16_t called;
+    psa_key_slot_number_t slot_number;
+    size_t data_size;
+    psa_status_t return_value;
+} mock_export_data;
+
+static struct
+{
+    uint16_t called;
+    psa_status_t return_value;
+} mock_allocate_data;
+
+static struct
+{
+    uint16_t called;
+    psa_key_slot_number_t slot_number;
+    psa_status_t return_value;
+} mock_destroy_data;
+
+static void mock_teardown( void )
+{
+    memset( &mock_import_data, 0, sizeof( mock_import_data ) );
+    memset( &mock_export_data, 0, sizeof( mock_export_data ) );
+    memset( &mock_allocate_data, 0, sizeof( mock_allocate_data ) );
+    memset( &mock_destroy_data, 0, sizeof( mock_destroy_data ) );
+    memset( &mock_generate_data, 0, sizeof( mock_generate_data ) );
+}
+
+static psa_status_t mock_generate( psa_key_slot_number_t key_slot,
+                                   psa_key_type_t type,
+                                   psa_key_usage_t usage,
+                                   size_t bits,
+                                   const void *extra,
+                                   size_t extra_size,
+                                   uint8_t *p_pubkey_out,
+                                   size_t pubkey_out_size,
+                                   size_t *p_pubkey_length )
+{
+    (void) extra;
+    (void) p_pubkey_out;
+    (void) p_pubkey_length;
+
+    mock_generate_data.called++;
+    mock_generate_data.key_slot = key_slot;
+    mock_generate_data.type = type;
+    mock_generate_data.usage = usage;
+    mock_generate_data.bits = bits;
+    mock_generate_data.extra_size = extra_size;
+    mock_generate_data.pubkey_out_size = pubkey_out_size;
+
+    return( mock_generate_data.return_value );
+}
+
+static psa_status_t mock_import( psa_key_slot_number_t slot_number,
+                                 psa_key_lifetime_t lifetime,
+                                 psa_key_type_t type,
+                                 psa_algorithm_t algorithm,
+                                 psa_key_usage_t usage,
+                                 const uint8_t *p_data,
+                                 size_t data_length )
+{
+    (void) p_data;
+
+    mock_import_data.called++;
+    mock_import_data.slot_number = slot_number;
+    mock_import_data.lifetime = lifetime;
+    mock_import_data.type = type;
+    mock_import_data.algorithm = algorithm;
+    mock_import_data.usage = usage;
+    mock_import_data.data_length = data_length;
+
+    return( mock_import_data.return_value );
+}
+
+psa_status_t mock_export( psa_key_slot_number_t slot_number,
+                          uint8_t *p_data,
+                          size_t data_size,
+                          size_t *p_data_length )
+{
+    (void) p_data;
+    (void) p_data_length;
+
+    mock_export_data.called++;
+    mock_export_data.slot_number = slot_number;
+    mock_export_data.data_size = data_size;
+
+    return( mock_export_data.return_value );
+}
+
+psa_status_t mock_allocate( const psa_key_attributes_t *attributes,
+                           const psa_drv_se_slot_usage_t *slot_usage,
+                           psa_key_slot_number_t *slot_number )
+{
+    (void) attributes;
+    (void) slot_usage;
+    (void) slot_number;
+
+    mock_allocate_data.called++;
+    *slot_number = 0;
+
+    return( mock_allocate_data.return_value );
+}
+
+psa_status_t mock_destroy( psa_key_slot_number_t slot_number )
+{
+    mock_destroy_data.called++;
+    mock_destroy_data.slot_number = slot_number;
+
+    return( mock_destroy_data.return_value );
+}
+
+/* END_HEADER */
+
+/* BEGIN_DEPENDENCIES
+ * depends_on:MBEDTLS_PSA_CRYPTO_SE_C
+ * END_DEPENDENCIES
+ */
+
+/* BEGIN_CASE */
+void mock_import( int mock_alloc_return_value,
+                  int mock_import_return_value,
+                  int expected_result)
+{
+    psa_drv_se_t driver;
+    psa_drv_se_key_management_t key_management;
+    psa_key_lifetime_t lifetime = 2;
+    psa_key_id_t id = 1;
+    psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
+
+    mock_allocate_data.return_value = mock_alloc_return_value;
+    mock_import_data.return_value = mock_import_return_value;
+    memset( &driver, 0, sizeof( driver ) );
+    memset( &key_management, 0, sizeof( key_management ) );
+    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
+    driver.key_management = &key_management;
+    key_management.slot_count = 1;
+    key_management.p_import = mock_import;
+    key_management.p_destroy = mock_destroy;
+    key_management.p_allocate = mock_allocate;
+
+    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_id( &attributes, id );
+    psa_set_key_lifetime( &attributes, lifetime );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
+    TEST_ASSERT( psa_import_key( &attributes,
+                                key_material, sizeof( key_material ),
+                                &handle ) == expected_result );
+
+    TEST_ASSERT( mock_allocate_data.called == 1 );
+    TEST_ASSERT( mock_import_data.called ==
+                  (expected_result == PSA_SUCCESS? 1 : 0) );
+    TEST_ASSERT( mock_import_data.type ==
+                  (expected_result == PSA_SUCCESS? PSA_KEY_TYPE_RAW_DATA : 0));
+
+    if( expected_result == PSA_SUCCESS )
+    {
+        PSA_ASSERT( psa_destroy_key( handle ) );
+        TEST_ASSERT( mock_destroy_data.called == 1 );
+    }
+exit:
+    PSA_DONE( );
+    mock_teardown( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void mock_export( int mock_export_return_value, int expected_result )
+{
+    psa_drv_se_t driver;
+    psa_drv_se_key_management_t key_management;
+    psa_key_lifetime_t lifetime = 2;
+    psa_key_id_t id = 1;
+    psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
+    uint8_t exported[sizeof( key_material )];
+    size_t exported_length;
+
+    mock_export_data.return_value = mock_export_return_value;
+    memset( &driver, 0, sizeof( driver ) );
+    memset( &key_management, 0, sizeof( key_management ) );
+    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
+    driver.key_management = &key_management;
+    key_management.slot_count = 1;
+    key_management.p_import = mock_import;
+    key_management.p_export = mock_export;
+    key_management.p_destroy = mock_destroy;
+    key_management.p_allocate = mock_allocate;
+
+    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_id( &attributes, id );
+    psa_set_key_lifetime( &attributes, lifetime );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
+    PSA_ASSERT( psa_import_key( &attributes,
+                                key_material, sizeof( key_material ),
+                                &handle ) );
+
+    TEST_ASSERT( psa_export_key( handle,
+                                exported, sizeof( exported ),
+                                &exported_length ) == expected_result );
+
+    TEST_ASSERT( mock_export_data.called == 1 );
+
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
+    TEST_ASSERT( mock_destroy_data.called == 1 );
+
+exit:
+    PSA_DONE( );
+    mock_teardown( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void mock_generate( int mock_generate_return_value, int expected_result )
+{
+    psa_drv_se_t driver;
+    psa_drv_se_key_management_t key_management;
+    psa_key_lifetime_t lifetime = 2;
+    psa_key_id_t id = 1;
+    psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+    mock_generate_data.return_value = mock_generate_return_value;
+    memset( &driver, 0, sizeof( driver ) );
+    memset( &key_management, 0, sizeof( key_management ) );
+    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
+    driver.key_management = &key_management;
+    key_management.slot_count = 1;
+    key_management.p_generate = mock_generate;
+    key_management.p_destroy = mock_destroy;
+    key_management.p_allocate = mock_allocate;
+
+    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_id( &attributes, id );
+    psa_set_key_lifetime( &attributes, lifetime );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
+    TEST_ASSERT( psa_generate_key( &attributes, &handle ) == expected_result );
+    TEST_ASSERT( mock_generate_data.called == 1 );
+
+exit:
+    PSA_DONE( );
+    mock_teardown( );
+}
+/* END_CASE */