websocket: do unsubscribe before closing connection

Author: HaraldNordgren

graphql/websocket.go

@@ -194,16 +194,18 @@ func (w *webSocketClient) Close() error {
 	if w.conn == nil {
 		return nil
 	}
-	err := w.conn.WriteMessage(closeMessage, formatCloseMessage(closeNormalClosure, ""))
-	if err != nil {
-		return fmt.Errorf("failed to send closure message: %w", err)
-	}
-	err = w.UnsubscribeAll()
+	err := w.UnsubscribeAll()
 	if err != nil {
 		return fmt.Errorf("failed to unsubscribe: %w", err)
 	}
 	w.Lock()
 	defer w.Unlock()
+
+	err = w.conn.WriteMessage(closeMessage, formatCloseMessage(closeNormalClosure, ""))
+	if err != nil {
+		return fmt.Errorf("failed to send closure message: %w", err)
+	}
+
 	w.isClosing = true
 	close(w.errChan)
 	return w.conn.Close()

internal/integration/integration_test.go

@@ -242,6 +242,49 @@ func TestSubscriptionConnectionParams(t *testing.T) {
 	}
 }
 
+func TestSubscriptionClose(t *testing.T) {
+	_ = `# @genqlient
+	subscription count { count }`
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+
+	cases := []struct {
+		name  string
+		unsub bool
+	}{
+		{
+			name:  "unsubscribed_manually",
+			unsub: true,
+		},
+		{
+			name:  "unsubscribed_automatically",
+			unsub: false,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			wsClient := newRoundtripWebSocketClient(t, server.URL)
+
+			_, err := wsClient.Start(ctx)
+			require.NoError(t, err)
+
+			_, subscriptionID, err := count(ctx, wsClient)
+			require.NoError(t, err)
+
+			if tc.unsub {
+				err = wsClient.Unsubscribe(subscriptionID)
+				require.NoError(t, err)
+			}
+
+			err = wsClient.Close()
+			require.NoError(t, err)
+		})
+	}
+}
+
 func TestServerError(t *testing.T) {
 	_ = `# @genqlient
 	query failingQuery { fail me { id } }`