v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (2)

8349-8354: Prefer t.Cleanup/defer to guarantee body closure even on assertion failures

Closing the body is correct. To ensure it’s always closed even if the preceding require fails (or future changes add early returns), register the close with t.Cleanup right after the Do() succeeds.

-        resp, err := httpClient.Do(req)
-        require.NoError(t, err)
-        require.Equal(t, http.StatusOK, resp.StatusCode)
-        require.NoError(t, resp.Body.Close())
+        resp, err := httpClient.Do(req)
+        require.NoError(t, err)
+        t.Cleanup(func() { require.NoError(t, resp.Body.Close()) })
+        require.Equal(t, http.StatusOK, resp.StatusCode)

---

8477-8482: Mirror the cleanup pattern for consistent, robust resource handling

Same comment as above: use t.Cleanup (or defer) so the response body is closed even if an assertion aborts the test.

-        resp, err := httpClient.Do(req)
-        require.NoError(t, err)
-        require.Equal(t, http.StatusOK, resp.StatusCode)
-        require.NoError(t, resp.Body.Close())
+        resp, err := httpClient.Do(req)
+        require.NoError(t, err)
+        t.Cleanup(func() { require.NoError(t, resp.Body.Close()) })
+        require.Equal(t, http.StatusOK, resp.StatusCode)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these settings in your CodeRabbit configuration.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go (1)

489-548: Fix deferred-close guard: connection can leak on invalid subprotocol return

The deferred close relies on the local variable err to be non-nil. However, the default branch in the subprotocol switch returns a new error literal without assigning it to err. In that path, the deferred close won't run and the TCP connection leaks.

Concretely, the default branch around Line 547 returns without touching err. Assign the error to err before returning so the defer fires.

Apply this diff to make the deferred close effective for that path:

 switch subProtocol {
 case ProtocolGraphQLWS:
   return newGQLWSConnectionHandler(requestContext, engineContext, conn, options, updater, c.log), nil
 case ProtocolGraphQLTWS:
   return newGQLTWSConnectionHandler(requestContext, engineContext, conn, options, updater, c.log), nil
 default:
-  return nil, NewInvalidWsSubprotocolError(subProtocol)
+  err = NewInvalidWsSubprotocolError(subProtocol)
+  return nil, err
 }

Optional (safer): Convert the function to use a named return error and check that in the defer to cover all return sites without manual assignments.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go (2)

591-597: Reconsider closing Transport-owned net.Conn on Do error

When httpClient.Do(req) returns an error, net/http’s Transport is responsible for connection lifecycle. Closing info.Conn yourself can fight the Transport (e.g., racing with its own teardown or preventing reuse in some edge cases). It’s generally safer to let Transport manage it.

Apply this diff to leave connection management to net/http:

 upgradeResponse, err := c.httpClient.Do(req)
 if err != nil {
-  if conn != nil {
-    conn.Close()
-  }
   return nil, "", err
 }

If you’ve observed a concrete resource leak without this manual Close, please share details; otherwise I recommend relying on Transport here.

---

599-606: Optionally drain non-101 response before closing to enable connection reuse

For non-upgrade responses, draining the body before Close helps net/http reuse the TCP connection. This is optional but can improve performance for repeated attempts.

Apply this diff:

 // On failed upgrade, we close the connection here without transferring ownership to the callee.
 if upgradeResponse.StatusCode != http.StatusSwitchingProtocols {
-  upgradeResponse.Body.Close()
+  // Drain to EOF to allow connection reuse by net/http.
+  _, _ = io.Copy(io.Discard, upgradeResponse.Body)
+  upgradeResponse.Body.Close()
   return nil, "", &UpgradeRequestError{
     URL:        u,
     StatusCode: upgradeResponse.StatusCode,
   }
 }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these settings in your CodeRabbit configuration.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go (1)

509-527: initMsg rename avoids shadowing; add a copy to avoid aliasing the global buffer

getConnectionInitMessage can return the global connectionInitMessage slice. Mutating initMsg via jsonparser.Set could share the underlying array if capacity permits. Safer to copy before any Set calls.

Apply this minimal copy after retrieving initMsg:

 initMsg, err := c.getConnectionInitMessage(requestContext, options.URL, options.Header)
 if err != nil {
   return nil, err
 }
 
+// Ensure we don't accidentally mutate a shared/global backing array.
+initMsg = append([]byte(nil), initMsg...)
 
 if len(options.InitialPayload) > 0 {
   initMsg, err = jsonparser.Set(initMsg, options.InitialPayload, "payload")
   if err != nil {
     return nil, err
   }
 }
 
 if options.Body.Extensions != nil {
   initMsg, err = jsonparser.Set(initMsg, options.Body.Extensions, "payload", "extensions")
   if err != nil {
     return nil, err
   }
 }

Alternatively, modify getConnectionInitMessage to always return a fresh copy when no payload is supplied:

func (c *subscriptionClient) getConnectionInitMessage(ctx context.Context, url string, header http.Header) ([]byte, error) {
	if c.onWsConnectionInitCallback == nil {
		return append([]byte(nil), connectionInitMessage...), nil
	}
	// ... and likewise when len(payload) == 0
}

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

2348-2461: Great coverage of upgrade failure matrix; consider asserting URL on error (nit)

Table-driven tests cover a wide range of non-101 responses and assert the custom error type. As a minor enhancement, also assert that UpgradeRequestError.URL is set to the HTTP URL used internally (http://...) to lock in behavior.

Example assertion (after ErrorAs):

• require.Contains(t, upgradeErr.URL, "http://")

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

1903-1971: Avoid closing signaling channels; use buffered + non-blocking send to prevent panics on multiple pings

In the server loop, if the client sends more than one ping before the client’s complete arrives, closing pingReceived again will panic ("close of closed channel"). Same for payloadSend if multiple sends happen. Prefer a buffered chan with a non-blocking send (or sync.Once) instead of close.

Apply this diff to make signaling idempotent and safe:

-        pingReceived := make(chan struct{}) // closed when the server receives a ping
-        payloadSend := make(chan struct{})  // closed when the server sends a payload
+        // buffered channels + non-blocking send to avoid double-close panics if events repeat
+        pingReceived := make(chan struct{}, 1) // signaled when the server receives a ping
+        payloadSend := make(chan struct{}, 1)  // signaled when the server sends a payload

And replace closing with non-blocking sends when signaling:

-                    receivedPing = true
-                    close(pingReceived)
+                    receivedPing = true
+                    select {
+                    case pingReceived <- struct{}{}:
+                    default:
+                    }
...
-                    close(payloadSend)
+                    select {
+                    case payloadSend <- struct{}{}:
+                    default:
+                    }

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

2021-2027: Nice: wait on payload signal instead of sleeping (addresses earlier flakiness)

Waiting for payloadSend before unsubscribing removes timing sensitivity and aligns with the previous review suggestion to replace sleeps with channels.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

2497-2503: "Correct length but wrong content" case isn't 20 bytes; craft a 20-byte value

Sec-WebSocket-Accept is a base64 of a 20-byte SHA-1. This case currently encodes a 27-byte string. It will still fail as invalid (so the test passes), but the case description doesn't match. Use a 20-byte value to reflect the intent.

Apply this diff:

-        {
-            name: "Correct length but wrong content",
-            acceptKeyHandler: func(challengeKey string) string {
-                return base64.StdEncoding.EncodeToString([]byte("wrong-content-here-20-bytes"))
-            },
-            expectError:   true,
-            errorContains: "invalid Sec-WebSocket-Accept",
-        },
+        {
+            name: "Correct length but wrong content",
+            acceptKeyHandler: func(challengeKey string) string {
+                // 20 bytes (not the SHA-1 of challengeKey+GUID), so length is correct but content is wrong
+                return base64.StdEncoding.EncodeToString([]byte("12345678901234567890"))
+            },
+            expectError:   true,
+            errorContains: "invalid Sec-WebSocket-Accept",
+        },

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

2208-2211: Fix double-close of serverDone (panic risk)

serverDone is closed in a defer; closing it again in the error branch will panic. Remove the inner close.

-        defer func() {
-            _ = conn.Close(websocket.StatusNormalClosure, "server done")
-            close(serverDone)
-        }()
+        defer func() {
+            _ = conn.Close(websocket.StatusNormalClosure, "server done")
+            close(serverDone)
+        }()
@@
-                // Signal that the server is done (connection closed)
-                close(serverDone)
-                return // Exit handler goroutine
+                // Connection closed; defer will close serverDone
+                return // Exit handler goroutine

Also applies to: 2256-2257

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (2)

2433-2476: Capture tc per iteration to future-proof subtests

Range-var capture can bite if these subtests become parallel later. Shadow tc inside the loop.

-    for _, tc := range testCases {
-        t.Run(tc.name, func(t *testing.T) {
+    for _, tc := range testCases {
+        tc := tc
+        t.Run(tc.name, func(t *testing.T) {

---

2514-2570: Also capture tc per iteration here

Same reasoning as above; add a per-iteration capture.

-    for _, tc := range testCases {
-        t.Run(tc.name, func(t *testing.T) {
+    for _, tc := range testCases {
+        tc := tc
+        t.Run(tc.name, func(t *testing.T) {

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client_test.go (1)

1947-1958: Avoid fixed-iteration read loop; rely on condition and context timeout instead

Limiting to 5 iterations risks missing a late message under load. Prefer looping until both conditions are met or the read context times out.

Apply this diff:

-            // Process up to 5 messages or until we see both a ping and a complete
-            for i := 0; i < 5; i++ {
-                if receivedPing && receivedComplete {
-                    break
-                }
-
-                _, data, err = conn.Read(readCtx)
+            // Process messages until we see both a ping and a complete or the read context times out
+            for !(receivedPing && receivedComplete) {
+                _, data, err = conn.Read(readCtx)
                 if err != nil {
                     // Connection closed or timeout
                     t.Logf("Connection read ended: %v", err)
                     break
                 }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.