From 88b64f94ecf305f7234fd25726e000b0c644c461 Mon Sep 17 00:00:00 2001
From: Igor Abdrakhimov <igorabd@amazon.com>
Date: Wed, 19 Feb 2025 15:13:17 -0800
Subject: [PATCH] Add test for mqtt3 websocket failed transform

---
 tests/CMakeLists.txt             |  2 ++
 tests/v3/connection_state_test.c | 37 ++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 67aeb1e7..95289f6e 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -108,6 +108,8 @@ add_test_case(mqtt_validation_failure_connect_invalid_client_id_utf8)
 add_test_case(mqtt_validation_failure_invalid_will_topic_utf8)
 add_test_case(mqtt_validation_failure_invalid_username_utf8)
 
+add_test_case(mqtt_websocket_failed_transform)
+
 # Operation statistics tests
 add_test_case(mqtt_operation_statistics_simple_publish)
 add_test_case(mqtt_operation_statistics_offline_publish)
diff --git a/tests/v3/connection_state_test.c b/tests/v3/connection_state_test.c
index d895b44a..989d9307 100644
--- a/tests/v3/connection_state_test.c
+++ b/tests/v3/connection_state_test.c
@@ -3968,3 +3968,40 @@ AWS_TEST_CASE_FIXTURE(
     s_test_mqtt_validation_failure_invalid_username_utf8_fn,
     s_clean_up_mqtt_server_fn,
     &test_data)
+
+static void s_mqtt_client_test_websocket_failed_transform(
+    struct aws_http_message *request,
+    void *user_data,
+    aws_mqtt_transform_websocket_handshake_complete_fn *complete_fn,
+    void *complete_ctx) {
+
+    (void)user_data;
+
+    (*complete_fn)(request, AWS_ERROR_INVALID_STATE, complete_ctx);
+}
+
+static int s_test_mqtt_websocket_failed_transform_fn(struct aws_allocator *allocator, void *ctx) {
+    (void)allocator;
+    struct mqtt_connection_state_test *state_test_data = ctx;
+
+    struct aws_mqtt_connection_options connection_options = {
+        .socket_options = &state_test_data->socket_options,
+    };
+
+    aws_mqtt_client_connection_use_websockets(
+        state_test_data->mqtt_connection, s_mqtt_client_test_websocket_failed_transform, NULL, NULL, NULL);
+
+    ASSERT_SUCCESS(aws_mqtt_client_connection_connect(state_test_data->mqtt_connection, &connection_options));
+    aws_test311_wait_for_connection_to_fail(state_test_data);
+
+    ASSERT_INT_EQUALS(AWS_ERROR_INVALID_STATE, state_test_data->error);
+
+    return AWS_OP_SUCCESS;
+}
+
+AWS_TEST_CASE_FIXTURE(
+    mqtt_websocket_failed_transform,
+    s_setup_mqtt_server_fn,
+    s_test_mqtt_websocket_failed_transform_fn,
+    s_clean_up_mqtt_server_fn,
+    &test_data)