diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 2e1c626a65..0f376efe0f 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -3118,6 +3118,12 @@ void pake_operations(data_t *pw_data, int forced_status_setup_arg, int forced_st
                                       input_buffer, size_key_share),
                        PSA_SUCCESS);
 
+            /* Simulate that we are ready to get implicit key. */
+            operation.computation_stage.data.jpake_computation_stage.input_step =
+                PSA_PAKE_STEP_DERIVE;
+            operation.computation_stage.data.jpake_computation_stage.output_step =
+                PSA_PAKE_STEP_DERIVE;
+
             /* --- psa_pake_get_implicit_key --- */
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;
             mbedtls_test_driver_pake_hooks.hits = 0;
@@ -3242,9 +3248,22 @@ void ecjpake_rounds(int alg_arg, int primitive_arg, int hash_arg,
     ecjpake_do_round(alg, primitive_arg, &server, &client,
                      client_input_first, 2);
 
+    /* After get the key is obtained operation is aborted.
+       Adapt counter of expected hits. */
+    if (pake_in_driver) {
+        pake_expected_hit_count++;
+    }
+
     PSA_ASSERT(psa_pake_get_implicit_key(&server, &server_derive));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
                pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);
+
+    /* After get the key is obtained operation is aborted.
+       Adapt counter of expected hits. */
+    if (pake_in_driver) {
+        pake_expected_hit_count++;
+    }
+
     PSA_ASSERT(psa_pake_get_implicit_key(&client, &client_derive));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits,
                pake_in_driver ? pake_expected_hit_count++ : pake_expected_hit_count);