test timing

Author: HaraldNordgren

internal/integration/integration_test.go

@@ -252,10 +252,6 @@ func TestSubscriptionClose(t *testing.T) {
 	_ = `# @genqlient
 	subscription countClose { countClose }`
 
-	ctx := context.Background()
-	server := server.RunServer()
-	defer server.Close()
-
 	cases := []struct {
 		name  string
 		unsub bool
@@ -272,6 +268,12 @@ func TestSubscriptionClose(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
+			server := server.RunServer()
+			defer server.Close()
+
 			wsClient := newRoundtripWebSocketClient(t, server.URL)
 
 			_, err := wsClient.Start(ctx)

internal/integration/server/server.go

@@ -157,32 +157,38 @@ func (m mutationResolver) CreateUser(ctx context.Context, input NewUser) (*User,
 	return &newUser, nil
 }
 
-func countTo(stopCount int) (<-chan int, error) {
+func countTo(ctx context.Context, stopCount int) (<-chan int, error) {
 	respChan := make(chan int, 1)
 	go func(respChan chan int) {
 		defer close(respChan)
 		for counter := range stopCount {
-			respChan <- counter
-			time.Sleep(100 * time.Millisecond)
+			select {
+			case <-ctx.Done():
+				fmt.Println("ctx done:", ctx.Err())
+				return
+			default:
+				respChan <- counter
+				time.Sleep(100 * time.Millisecond)
+			}
 		}
 	}(respChan)
 	return respChan, nil
 }
 
 func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
-	return countTo(10)
+	return countTo(ctx, 10)
 }
 
 func (s *subscriptionResolver) CountAuthorized(ctx context.Context) (<-chan int, error) {
 	if getAuthToken(ctx) != "authorized-user-token" {
 		return nil, fmt.Errorf("unauthorized")
 	}
 
-	return countTo(10)
+	return countTo(ctx, 10)
 }
 
 func (s *subscriptionResolver) CountClose(ctx context.Context) (<-chan int, error) {
-	return countTo(1000)
+	return countTo(ctx, 1000)
 }
 
 const AuthKey = "authToken"