Function TestInMemoryRoundTripper_CancelCloses

coderd/aibridged/transport_test.go:191–226

Canceling the request context must surface as a body-read error, matching real-network behavior, and the handler must observe the cancellation through its own request context.

(t *testing.T)

// real-network behavior, and the handler must observe the cancellation
// through its own request context.
func TestInMemoryRoundTripper_CancelCloses(t *testing.T) {
	t.Parallel()

	handlerCtxObserved := make(chan struct{})
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
		<-r.Context().Done()
		close(handlerCtxObserved)
	})

	rt, err := aibridged.NewTransportFactory(handler).TransportFor("openai", aibridge.SourceAgents)
	require.NoError(t, err)

	parentCtx := testutil.Context(t, testutil.WaitShort)
	ctx, cancel := context.WithCancel(parentCtx)
	ctx = aibridge.WithDelegatedAPIKeyID(ctx, "test-key-id")
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://aibridge/stream", nil)
	require.NoError(t, err)

	resp, err := rt.RoundTrip(req)
	require.NoError(t, err)
	defer resp.Body.Close()

	cancel()
	_, err = io.ReadAll(resp.Body)
	require.Error(t, err)

	select {
	case <-handlerCtxObserved:
	case <-parentCtx.Done():
		t.Fatal("handler did not observe context cancellation")
	}
}

// Many independent in-flight requests on a shared handler must not interfere.
func TestInMemoryRoundTripper_ConcurrentRequests(t *testing.T) {