diff --git a/CMakeLists.txt b/CMakeLists.txt
index a3cf1ca38ac..d1665df1304 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -198,11 +198,10 @@ if(ASAN)
 endif()
 
 if (UBSAN)
-    target_compile_options(${PROJECT_NAME} PUBLIC -fsanitize=undefined -fno-sanitize-recover=all -DS2N_ADDRESS_SANITIZER=1)
+    target_compile_options(${PROJECT_NAME} PUBLIC -fsanitize=undefined -fno-sanitize-recover=all)
     target_link_options(${PROJECT_NAME} PUBLIC -fsanitize=undefined -fno-sanitize-recover=all)
 endif()
 
-
 if(TSAN OR ASAN OR UBSAN)
     # no-omit-frame-pointer and no-optimize-sibling-calls provide better stack traces
     target_compile_options(${PROJECT_NAME} PUBLIC -fno-omit-frame-pointer -fno-optimize-sibling-calls)
diff --git a/tests/testlib/s2n_key_schedule_testlib.c b/tests/testlib/s2n_key_schedule_testlib.c
index a539120eb08..fab73425309 100644
--- a/tests/testlib/s2n_key_schedule_testlib.c
+++ b/tests/testlib/s2n_key_schedule_testlib.c
@@ -20,7 +20,6 @@ S2N_RESULT s2n_connection_set_test_transcript_hash(struct s2n_connection *conn,
 {
     RESULT_GUARD(s2n_connection_set_test_message_type(conn, message_type));
     conn->handshake.handshake_type &= NEGOTIATED;
-    RESULT_ENSURE_EQ(s2n_conn_get_current_message_type(conn), message_type);
     RESULT_CHECKED_MEMCPY(conn->handshake.hashes->transcript_hash_digest,
             digest->data, digest->size);
     return S2N_RESULT_OK;
@@ -59,9 +58,14 @@ S2N_RESULT s2n_connection_set_test_master_secret(struct s2n_connection *conn,
     return S2N_RESULT_OK;
 }
 
+/* This function will iterate over all rows and columns of the handshake state 
+ * machine until it finds a valid (handshake_type, handshake_number) such that 
+ * the active message is `expected_message_type`. If callers need to depend on a
+ * specific `message_number` or `handshake_type` this function should not be 
+ * used.
+ */
 S2N_RESULT s2n_connection_set_test_message_type(struct s2n_connection *conn, message_type_t expected_message_type)
 {
-    /* scan over all handshake types until the desired message type is found */
     for (uint32_t handshake = 0; handshake < S2N_HANDSHAKES_COUNT; handshake++) {
         for (int message = 0; message < S2N_MAX_HANDSHAKE_LENGTH; message++) {
             conn->handshake.handshake_type = handshake;