chore: Fix flaky unit test (#4374)

Author: gmlewis

github/github_test.go

@@ -2846,6 +2846,9 @@ func TestDo_rateLimit_abortSleepContextCancelled(t *testing.T) {
 	// We use a 1 minute reset time to ensure the sleep is not completed.
 	reset := time.Now().UTC().Add(time.Minute)
 	var requestCount atomic.Int32
+	// requestReceived is signaled once the handler has run, so the test can
+	// verify the request count independently of when the context is cancelled.
+	requestReceived := make(chan struct{})
 	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
 		requestCount.Add(1)
 		w.Header().Set(HeaderRateLimit, "60")
@@ -2859,18 +2862,46 @@ func TestDo_rateLimit_abortSleepContextCancelled(t *testing.T) {
    "message": "API rate limit exceeded for xxx.xxx.xxx.xxx. (But here's the good news: Authenticated requests get a higher rate limit. Check out the documentation for more details.)",
    "documentation_url": "https://docs.github.com/en/rest/overview/resources-in-the-rest-api#abuse-rate-limits"
 }`)
+		close(requestReceived)
 	})
 
-	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Millisecond)
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
 	req, _ := client.NewRequest(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), "GET", ".", nil)
-	_, err := client.Do(req, nil)
-	if !errors.Is(err, context.DeadlineExceeded) {
-		t.Error("Expected context deadline exceeded error.")
+
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := client.Do(req, nil)
+		errCh <- err
+	}()
+
+	// Wait until the handler has processed the request, ensuring exactly one
+	// request was made before we cancel the context to abort the rate-limit
+	// sleep. This decouples the request-count assertion from the timing of
+	// the context cancellation, eliminating flakiness on slow runners.
+	select {
+	case <-requestReceived:
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timed out waiting for request to be received.")
 	}
+
 	if got, want := int(requestCount.Load()), 1; got != want {
-		t.Errorf("Expected 1 requests, got %v", got)
+		t.Errorf("Expected 1 request, got %v", got)
+	}
+
+	// Now cancel the context while the client is sleeping waiting for the
+	// rate limit to reset. This should abort the sleep and return the
+	// context error.
+	cancel()
+
+	select {
+	case err := <-errCh:
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("Expected context cancelled error, got: %v", err)
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timed out waiting for Do to return after context cancellation.")
 	}
 }