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
index c05b182..dba6875 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.data
@@ -1,11 +1,17 @@
 SE key importing mock test
-mock_import:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+mock_import:PSA_SUCCESS:PSA_SUCCESS:0:PSA_SUCCESS
+
+SE key importing mock test: max key bits
+mock_import:PSA_SUCCESS:PSA_SUCCESS:PSA_MAX_KEY_BITS:PSA_SUCCESS
+
+SE key importing mock test: more than max key bits
+mock_import:PSA_SUCCESS:PSA_ERROR_NOT_SUPPORTED:PSA_MAX_KEY_BITS+1:PSA_ERROR_NOT_SUPPORTED
 
 SE key importing mock test: alloc failed
-mock_import:PSA_ERROR_HARDWARE_FAILURE:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE
+mock_import:PSA_ERROR_HARDWARE_FAILURE:PSA_SUCCESS:0:PSA_ERROR_HARDWARE_FAILURE
 
 SE key importing mock test: import failed
-mock_import:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:PSA_ERROR_HARDWARE_FAILURE
+mock_import:PSA_SUCCESS:PSA_ERROR_HARDWARE_FAILURE:0:PSA_ERROR_HARDWARE_FAILURE
 
 SE key exporting mock test
 mock_export:PSA_SUCCESS:PSA_SUCCESS
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
index 9d73d8f..e364178 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
@@ -19,6 +19,7 @@
     uint16_t called;
     psa_key_slot_number_t key_slot;
     psa_key_attributes_t attributes;
+    size_t bits;
     size_t data_length;
     psa_status_t return_value;
 } mock_import_data;
@@ -130,7 +131,8 @@
 {
     (void) drv_context;
     (void) data;
-    (void) bits;
+
+    *bits = mock_import_data.bits;
 
     mock_import_data.called++;
     mock_import_data.key_slot = key_slot;
@@ -259,6 +261,7 @@
 /* BEGIN_CASE */
 void mock_import( int mock_alloc_return_value,
                   int mock_import_return_value,
+                  int bits,
                   int expected_result )
 {
     psa_drv_se_t driver;
@@ -271,6 +274,7 @@
 
     mock_allocate_data.return_value = mock_alloc_return_value;
     mock_import_data.return_value = mock_import_return_value;
+    mock_import_data.bits = bits;
     memset( &driver, 0, sizeof( driver ) );
     memset( &key_management, 0, sizeof( key_management ) );
     driver.hal_version = PSA_DRV_SE_HAL_VERSION;