Dispatch according to init status.

We shouldn't dispatch to PSA when drivers have not been initialized yet.

Signed-off-by: Manuel Pégourié-Gonnard <manuel.pegourie-gonnard@arm.com>
diff --git a/tests/suites/test_suite_md.function b/tests/suites/test_suite_md.function
index 1e8622b..bf875f3 100644
--- a/tests/suites/test_suite_md.function
+++ b/tests/suites/test_suite_md.function
@@ -16,6 +16,9 @@
     unsigned char out[MBEDTLS_MD_MAX_SIZE] = { 0 };
 
     mbedtls_md_init(&ctx);
+#if defined(MBEDTLS_MD_SOME_PSA)
+    PSA_INIT();
+#endif
 
     /*
      * Test that mbedtls_md_list() only returns valid MDs.
@@ -31,6 +34,9 @@
 
 exit:
     mbedtls_md_free(&ctx);
+#if defined(MBEDTLS_MD_SOME_PSA)
+    PSA_DONE();
+#endif
 }
 /* END_CASE */
 
@@ -338,3 +344,48 @@
     ASSERT_COMPARE(output, mbedtls_md_get_size(md_info), hash->x, hash->len);
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void md_psa_dynamic_dispatch(int md_type, int pre_psa_ret, int post_psa_engine)
+{
+    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type(md_type);
+    TEST_ASSERT(md_info != NULL);
+    mbedtls_md_context_t ctx1, ctx2;
+
+    mbedtls_md_init(&ctx1);
+    mbedtls_md_init(&ctx2);
+
+    /* Before PSA crypto init */
+    TEST_EQUAL(pre_psa_ret, mbedtls_md_setup(&ctx1, md_info, 0));
+    TEST_EQUAL(pre_psa_ret, mbedtls_md_setup(&ctx2, md_info, 0));
+
+#if defined(MBEDTLS_MD_SOME_PSA)
+    TEST_EQUAL(ctx1.engine, MBEDTLS_MD_ENGINE_LEGACY);
+    TEST_EQUAL(ctx2.engine, MBEDTLS_MD_ENGINE_LEGACY);
+#endif
+
+    /* Reset ctx1 but keep ctx2 for the cloning test */
+    mbedtls_md_free(&ctx1);
+    mbedtls_md_init(&ctx1);
+
+    /* After PSA Crypto init */
+    PSA_INIT();
+    TEST_EQUAL(0, mbedtls_md_setup(&ctx1, md_info, 0));
+#if defined(MBEDTLS_MD_SOME_PSA)
+    TEST_EQUAL(ctx1.engine, post_psa_engine);
+#endif
+
+    /* Cloning test */
+    if (pre_psa_ret == 0) {
+        int exp_clone_ret = post_psa_engine == MBEDTLS_MD_ENGINE_PSA
+                          ? MBEDTLS_ERR_MD_FEATURE_UNAVAILABLE
+                          : 0;
+        TEST_EQUAL(exp_clone_ret, mbedtls_md_clone(&ctx2, &ctx1));
+    }
+
+exit:
+    mbedtls_md_free(&ctx1);
+    mbedtls_md_free(&ctx2);
+    PSA_DONE();
+}
+/* END_CASE */