diff --git a/library/pk.c b/library/pk.c
index d8bce8f..38ab774 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -90,8 +90,7 @@
  */
 void mbedtls_pk_restart_free( mbedtls_pk_restart_ctx *ctx )
 {
-    PK_VALIDATE( ctx != NULL );
-    if( ctx->pk_info == NULL ||
+    if( ctx == NULL || ctx->pk_info == NULL ||
         ctx->pk_info->rs_free_func == NULL )
     {
         return;
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 936bcdf..b8069b5 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -77,6 +77,12 @@
 
     mbedtls_pk_init( &pk );
 
+    TEST_VALID_PARAM( mbedtls_pk_free( NULL ) );
+
+#if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
+    TEST_VALID_PARAM( mbedtls_pk_restart_free( NULL ) );
+#endif
+
     TEST_ASSERT( mbedtls_pk_setup( &pk, NULL ) ==
                  MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
@@ -133,8 +139,6 @@
 
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     TEST_INVALID_PARAM( mbedtls_pk_restart_init( NULL ) );
-
-    TEST_INVALID_PARAM( mbedtls_pk_restart_free( NULL ) );
 #endif
 
     TEST_INVALID_PARAM_RET( MBEDTLS_ERR_PK_BAD_INPUT_DATA,