Use correct export length in LMS export tests

Signed-off-by: Raef Coles <raef.coles@arm.com>
diff --git a/tests/suites/test_suite_lmots.function b/tests/suites/test_suite_lmots.function
index b5b18b5..53ab02f 100644
--- a/tests/suites/test_suite_lmots.function
+++ b/tests/suites/test_suite_lmots.function
@@ -182,10 +182,9 @@
         exported_pub_key_buf_size = MBEDTLS_LMOTS_PUBLIC_KEY_LEN(MBEDTLS_LMOTS_SHA256_N32_W8) + 1;
         ASSERT_ALLOC( exported_pub_key, exported_pub_key_buf_size);
         TEST_EQUAL( mbedtls_lmots_export_public_key( &ctx, exported_pub_key,
-                                                   exported_pub_key_buf_size, NULL ),
+                                                   exported_pub_key_buf_size,
+                                                   &exported_pub_key_size ),
                     0 );
-        TEST_EQUAL( exported_pub_key_size,
-                    MBEDTLS_LMOTS_PUBLIC_KEY_LEN(MBEDTLS_LMOTS_SHA256_N32_W8) );
         ASSERT_COMPARE( pub_key->x, pub_key->len,
                         exported_pub_key, exported_pub_key_size );
         mbedtls_free(exported_pub_key);
diff --git a/tests/suites/test_suite_lms.function b/tests/suites/test_suite_lms.function
index baf74ab..c5c8aa4 100644
--- a/tests/suites/test_suite_lms.function
+++ b/tests/suites/test_suite_lms.function
@@ -184,10 +184,9 @@
         exported_pub_key_buf_size = MBEDTLS_LMS_PUBLIC_KEY_LEN(MBEDTLS_LMS_SHA256_M32_H10) + 1;
         ASSERT_ALLOC( exported_pub_key, exported_pub_key_buf_size);
         TEST_EQUAL( mbedtls_lms_export_public_key( &ctx, exported_pub_key,
-                                                   exported_pub_key_buf_size, NULL ),
+                                                   exported_pub_key_buf_size,
+                                                   &exported_pub_key_size ),
                     0 );
-        TEST_EQUAL( exported_pub_key_size,
-                    MBEDTLS_LMS_PUBLIC_KEY_LEN(MBEDTLS_LMS_SHA256_M32_H10 ) );
         ASSERT_COMPARE( pub_key->x, pub_key->len,
                         exported_pub_key, exported_pub_key_size );
         mbedtls_free(exported_pub_key);