Create separate test function

internal/integration/integration_test.go

@@ -66,17 +66,13 @@ func TestSubscription(t *testing.T) {
 	_ = `# @genqlient
 	subscription count { count }`
 
-	authKey := server.AuthKey
-
 	ctx := context.Background()
 	server := server.RunServer()
 	defer server.Close()
 
 	cases := []struct {
-		connParams     map[string]interface{}
 		name           string
 		unsubThreshold time.Duration
-		counterStart   int
 		expected       subscriptionResult
 	}{
 		{
@@ -87,18 +83,6 @@ func TestSubscription(t *testing.T) {
 				serverChannelClosed: true,
 			},
 		},
-		{
-			name:           "server_closed_authorized_user_gets_incremented_counter",
-			unsubThreshold: 5 * time.Second,
-			counterStart:   1000,
-			connParams: map[string]interface{}{
-				authKey: "authorized-user-token",
-			},
-			expected: subscriptionResult{
-				clientUnsubscribed:  false,
-				serverChannelClosed: true,
-			},
-		},
 		{
 			name:           "client_unsubscribed",
 			unsubThreshold: 300 * time.Millisecond,
@@ -111,7 +95,7 @@ func TestSubscription(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			wsClient := newRoundtripWebSocketClient(t, server.URL, tc.connParams)
+			wsClient := newRoundtripWebSocketClient(t, server.URL, nil)
 
 			errChan, err := wsClient.Start(ctx)
 			require.NoError(t, err)
@@ -121,7 +105,7 @@ func TestSubscription(t *testing.T) {
 			defer wsClient.Close()
 
 			var (
-				counter = tc.counterStart
+				counter = 0
 				start   = time.Now()
 				result  = subscriptionResult{}
 			)
@@ -161,6 +145,86 @@ func TestSubscription(t *testing.T) {
 	}
 }
 
+func TestSubscriptionConnectionParams(t *testing.T) {
+	_ = `# @genqlient
+	subscription countAuthorized { countAuthorized }`
+
+	authKey := server.AuthKey
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+
+	cases := []struct {
+		connParams    map[string]interface{}
+		name          string
+		expectedError string
+	}{
+		{
+			name: "authorized_user_gets_counter",
+			connParams: map[string]interface{}{
+				authKey: "authorized-user-token",
+			},
+		},
+		{
+			name:          "unauthorized_user_gets_error",
+			expectedError: "input: countAuthorized unauthorized\n",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			wsClient := newRoundtripWebSocketClient(t, server.URL, tc.connParams)
+
+			errChan, err := wsClient.Start(ctx)
+			require.NoError(t, err)
+
+			dataChan, subscriptionID, err := countAuthorized(ctx, wsClient)
+			require.NoError(t, err)
+			defer wsClient.Close()
+
+			var (
+				counter = 0
+				start   = time.Now()
+			)
+
+			for loop := true; loop; {
+				select {
+				case resp, more := <-dataChan:
+					if !more {
+						loop = false
+						break
+					}
+
+					if tc.expectedError != "" {
+						require.Error(t, resp.Errors)
+						assert.Equal(t, tc.expectedError, resp.Errors.Error())
+						continue
+					}
+
+					require.NotNil(t, resp.Data)
+					assert.Equal(t, counter, resp.Data.CountAuthorized)
+					require.Nil(t, resp.Errors)
+
+					if time.Since(start) > 5*time.Second {
+						err := wsClient.Unsubscribe(subscriptionID)
+						require.NoError(t, err)
+						loop = false
+					}
+
+					counter++
+
+				case err := <-errChan:
+					require.NoError(t, err)
+
+				case <-time.After(10 * time.Second):
+					require.NoError(t, fmt.Errorf("subscription timed out"))
+				}
+			}
+		})
+	}
+}
+
 func TestServerError(t *testing.T) {
 	_ = `# @genqlient
 	query failingQuery { fail me { id } }`

internal/integration/schema.graphql

@@ -19,6 +19,7 @@ type Mutation {
 
 type Subscription {
   count: Int!
+  countAuthorized: Int!
 }
 
 type User implements Being & Lucky {

internal/integration/server/server.go

@@ -168,23 +168,37 @@ func getAuthToken(ctx context.Context) string {
 }
 
 func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
-	respCounter := 0
-	if getAuthToken(ctx) == "authorized-user-token" {
-		respCounter = 1000
+	respChan := make(chan int, 1)
+	go func(respChan chan int) {
+		defer close(respChan)
+		counter := 0
+		for {
+			if counter == 10 {
+				return
+			}
+			respChan <- counter
+			counter++
+			time.Sleep(100 * time.Millisecond)
+		}
+	}(respChan)
+	return respChan, nil
+}
+
+func (s *subscriptionResolver) CountAuthorized(ctx context.Context) (<-chan int, error) {
+	if getAuthToken(ctx) != "authorized-user-token" {
+		return nil, fmt.Errorf("unauthorized")
 	}
 
 	respChan := make(chan int, 1)
 	go func(respChan chan int) {
 		defer close(respChan)
-		closeCounter := 0
+		counter := 0
 		for {
-			if closeCounter == 10 {
+			if counter == 10 {
 				return
 			}
-			closeCounter++
-			respChan <- respCounter
-			respCounter++
-			time.Sleep(100 * time.Millisecond)
+			respChan <- counter
+			counter++
 		}
 	}(respChan)
 	return respChan, nil