Improve interuptible key generation testing

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
index 9d1ecef..4cd4eae 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
@@ -10122,6 +10122,11 @@
         TEST_EQUAL(mbedtls_test_psa_exercise_key(key, usage, alg, 0), 1);
     }
 
+    /* Adjust expected_status for interruptible key generation.
+     * Interruptible key generation is only supported for ECC key pairs and even
+     * for those only when MBEDTLS_ECP_RESTARTABLE is on.
+     */
+
     if (!PSA_KEY_TYPE_IS_ECC(type)) {
         expected_status = PSA_ERROR_NOT_SUPPORTED;
     }
@@ -10132,17 +10137,18 @@
 
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, expected_status);
-    if (expected_status != PSA_SUCCESS) {
-        goto exit;
-    }
 
+    /* Test that calling setup 2 time consecutively fails */
+#if defined(MBEDTLS_ECP_RESTARTABLE)
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+#endif
 
-    status = psa_generate_key_iop_abort(&operation);
+    PSA_ASSERT(psa_generate_key_iop_abort(&operation));
 
+    /* Test that after calling abort operation is reset to it's fresh state */
     status = psa_generate_key_iop_setup(&operation, &attributes);
-    TEST_EQUAL(status, PSA_SUCCESS);
+    TEST_EQUAL(status, expected_status);
 
 exit:
     psa_generate_key_iop_abort(&operation);