fix test fail when WANT_READ/WRITE returned

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 8a08963..6ba4e20 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -4910,12 +4910,16 @@
                                            state );
     if( need_pass )
     {
-        TEST_ASSERT( ret == 0 );
+        TEST_ASSERT( ret == 0 ||
+                     ret == MBEDTLS_ERR_SSL_WANT_READ ||
+                     ret == MBEDTLS_ERR_SSL_WANT_WRITE );
         TEST_ASSERT( base_ep.ssl.state == state );
     }
     else
     {
-        TEST_ASSERT( ret != 0 );
+        TEST_ASSERT( ret != 0 &&
+                     ret != MBEDTLS_ERR_SSL_WANT_READ &&
+                     ret != MBEDTLS_ERR_SSL_WANT_WRITE );
         TEST_ASSERT( base_ep.ssl.state != state );
     }