Fix _arg argument not being cast to correct type

Also change to TEST_EQUAL, as this is now possible.

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 0d9543d..5e4eaf8 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -3685,6 +3685,7 @@
     uint8_t nonce_buffer[PSA_AEAD_NONCE_MAX_SIZE];
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+    psa_status_t expected_status = expected_status_arg;
     size_t nonce_generated_len = 0;
     size_t expected_generated_len = expected_generated_len_arg;
     unsigned char *output_data = NULL;
@@ -3735,13 +3736,13 @@
                                       nonce_len,
                                       &nonce_generated_len );
 
-    TEST_ASSERT( status == expected_status_arg );
+    TEST_EQUAL( status, expected_status );
 
     TEST_EQUAL( nonce_generated_len, expected_generated_len );
 
     TEST_ASSERT( nonce_generated_len < PSA_AEAD_NONCE_MAX_SIZE );
 
-    if( expected_status_arg == PSA_SUCCESS )
+    if( expected_status == PSA_SUCCESS )
     {
 
         /* Ensure we can still complete operation. */
@@ -3837,7 +3838,7 @@
 
     status = psa_aead_set_nonce( &operation, nonce_buffer, nonce_len );
 
-    TEST_ASSERT( status == expected_status );
+    TEST_EQUAL( status, expected_status );
 
     if( expected_status == PSA_SUCCESS )
     {