diff --git a/unit_test/spdm_transport_test_lib/common.c b/unit_test/spdm_transport_test_lib/common.c
index 59604bfa2aa..314f5d5ef36 100644
--- a/unit_test/spdm_transport_test_lib/common.c
+++ b/unit_test/spdm_transport_test_lib/common.c
@@ -21,7 +21,9 @@
  * @retval RETURN_SUCCESS               The message is encoded successfully.
  * @retval RETURN_INVALID_PARAMETER     The message is NULL or the message_size is zero.
  **/
-libspdm_return_t libspdm_test_encode_message(const uint32_t *session_id, size_t message_size,
+libspdm_return_t libspdm_test_encode_message(const uint32_t *session_id,
+                                             bool need_alignment,
+                                             size_t message_size,
                                              const void *message,
                                              size_t *transport_message_size,
                                              void **transport_message);
@@ -116,7 +118,7 @@ libspdm_return_t libspdm_transport_test_encode_message(
             /* SPDM message to APP message*/
             app_message = NULL;
             app_message_size = transport_header_size + message_size + (LIBSPDM_TEST_ALIGNMENT - 1);
-            status = libspdm_test_encode_message(NULL, message_size,
+            status = libspdm_test_encode_message(NULL, false, message_size,
                                                  message,
                                                  &app_message_size,
                                                  &app_message);
@@ -145,7 +147,7 @@ libspdm_return_t libspdm_transport_test_encode_message(
 
         /* secured message to secured Test message*/
         status = libspdm_test_encode_message(
-            session_id, secured_message_size, secured_message,
+            session_id, true, secured_message_size, secured_message,
             transport_message_size, transport_message);
         if (LIBSPDM_STATUS_IS_ERROR(status)) {
             LIBSPDM_DEBUG((LIBSPDM_DEBUG_ERROR, "transport_encode_message - %xu\n",
@@ -154,7 +156,7 @@ libspdm_return_t libspdm_transport_test_encode_message(
         }
     } else {
         /* SPDM message to normal Test message*/
-        status = libspdm_test_encode_message(NULL, message_size, message,
+        status = libspdm_test_encode_message(NULL, true, message_size, message,
                                              transport_message_size,
                                              transport_message);
         if (LIBSPDM_STATUS_IS_ERROR(status)) {
diff --git a/unit_test/spdm_transport_test_lib/test.c b/unit_test/spdm_transport_test_lib/test.c
index 48d7cd66e8b..8a630ffc1eb 100644
--- a/unit_test/spdm_transport_test_lib/test.c
+++ b/unit_test/spdm_transport_test_lib/test.c
@@ -70,7 +70,9 @@ spdm_version_number_t libspdm_test_get_secured_spdm_version(
  * @retval RETURN_SUCCESS               The message is encoded successfully.
  * @retval RETURN_INVALID_PARAMETER     The message is NULL or the message_size is zero.
  **/
-libspdm_return_t libspdm_test_encode_message(const uint32_t *session_id, size_t message_size,
+libspdm_return_t libspdm_test_encode_message(const uint32_t *session_id,
+                                             bool need_alignment,
+                                             size_t message_size,
                                              void *message,
                                              size_t *transport_message_size,
                                              void **transport_message)
@@ -80,7 +82,11 @@ libspdm_return_t libspdm_test_encode_message(const uint32_t *session_id, size_t
     uint32_t data32;
     libspdm_test_message_header_t *test_message_header;
 
-    alignment = LIBSPDM_TEST_ALIGNMENT;
+    if (need_alignment) {
+        alignment = LIBSPDM_TEST_ALIGNMENT;
+    } else {
+        alignment = 1;
+    }
     aligned_message_size =
         (message_size + (alignment - 1)) & ~(alignment - 1);