Refactor

internal/integration/integration_test.go

@@ -57,45 +57,95 @@ func TestMutation(t *testing.T) {
 	require.Errorf(t, err, "client does not support mutations")
 }
 
-func TestSubscription(t *testing.T) {
-	_ = `# @genqlient
-	subscription count { count }`
-
-	ctx := context.Background()
-	server := server.RunServer()
-	defer server.Close()
-	wsClient := newRoundtripWebScoketClient(t, server.URL)
+type subscriptionCountLoopResult struct {
+	loop                     bool
+	clientTriggerUnsubscribe bool
+	serverClosedChannel      bool
+}
 
+func subscriptionCountLoop(
+	ctx context.Context,
+	t *testing.T,
+	wsClient graphql.WebSocketClient,
+	clientUnsubscribeDuration time.Duration,
+) subscriptionCountLoopResult {
 	errChan, err := wsClient.Start(ctx)
 	require.NoError(t, err)
 
 	dataChan, subscriptionID, err := count(ctx, wsClient)
 	require.NoError(t, err)
 	defer wsClient.Close()
-	counter := 0
-	start := time.Now()
-	for loop := true; loop; {
+
+	var (
+		counter                  = 0
+		start                    = time.Now()
+		loop                     = true
+		clientTriggerUnsubscribe = false
+		serverClosedChannel      = false
+	)
+
+	for loop {
 		select {
 		case resp, more := <-dataChan:
 			if !more {
 				loop = false
+				serverClosedChannel = true
 				break
 			}
 			require.NotNil(t, resp.Data)
 			assert.Equal(t, counter, resp.Data.Count)
 			require.Nil(t, resp.Errors)
-			if time.Since(start) > time.Second*5 {
-				err = wsClient.Unsubscribe(subscriptionID)
+			if time.Since(start) > clientUnsubscribeDuration {
+				err := wsClient.Unsubscribe(subscriptionID)
 				require.NoError(t, err)
 				loop = false
+				clientTriggerUnsubscribe = true
 			}
 			counter++
 		case err := <-errChan:
 			require.NoError(t, err)
-		case <-time.After(time.Second * 10):
+		case <-time.After(5 * time.Second):
 			require.NoError(t, fmt.Errorf("subscription timed out"))
 		}
 	}
+
+	return subscriptionCountLoopResult{
+		loop:                     loop,
+		clientTriggerUnsubscribe: clientTriggerUnsubscribe,
+		serverClosedChannel:      serverClosedChannel,
+	}
+}
+
+func TestSubscriptionServerClose(t *testing.T) {
+	_ = `# @genqlient
+	subscription count { count }`
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+	wsClient := newRoundtripWebScoketClient(t, server.URL)
+
+	result := subscriptionCountLoop(ctx, t, wsClient, 5*time.Second)
+
+	assert.False(t, result.loop)
+	assert.False(t, result.clientTriggerUnsubscribe)
+	assert.True(t, result.serverClosedChannel)
+}
+
+func TestSubscriptionClientClose(t *testing.T) {
+	_ = `# @genqlient
+	subscription count { count }`
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+	wsClient := newRoundtripWebScoketClient(t, server.URL)
+
+	result := subscriptionCountLoop(ctx, t, wsClient, 300*time.Millisecond)
+
+	assert.False(t, result.loop)
+	assert.True(t, result.clientTriggerUnsubscribe)
+	assert.False(t, result.serverClosedChannel)
 }
 
 func TestServerError(t *testing.T) {

internal/integration/server/server.go

@@ -162,12 +162,12 @@ func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
 		defer close(respChan)
 		counter := 0
 		for {
-			if counter == 3 {
+			if counter == 10 {
 				return
 			}
 			respChan <- counter
 			counter++
-			time.Sleep(10 * time.Millisecond)
+			time.Sleep(100 * time.Millisecond)
 		}
 	}(respChan)
 	return respChan, nil