From 8d0efb480599efe9ae40292164d7fdbf3d3e43fe Mon Sep 17 00:00:00 2001
From: Aman Shrivastava <amanshrivastava118@gmail.com>
Date: Tue, 23 Dec 2025 19:40:25 +0530
Subject: [PATCH] test: skip APF header assertions after request timeout

---
 .../filters/priority-and-fairness_test.go     | 53 +++++++++++--------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/staging/src/k8s.io/apiserver/pkg/server/filters/priority-and-fairness_test.go b/staging/src/k8s.io/apiserver/pkg/server/filters/priority-and-fairness_test.go
index 18d99292ccd..ad7105b477a 100644
--- a/staging/src/k8s.io/apiserver/pkg/server/filters/priority-and-fairness_test.go
+++ b/staging/src/k8s.io/apiserver/pkg/server/filters/priority-and-fairness_test.go
@@ -712,7 +712,7 @@ func TestPriorityAndFairnessWithPanicRecoveryAndTimeoutFilter(t *testing.T) {
 		firstRequestPathPanic, secondRequestPathShouldWork := "/request/panic-as-designed", "/request/should-succeed-as-expected"
 		firstHandlerDoneCh, secondHandlerDoneCh := make(chan struct{}), make(chan struct{})
 		requestHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			headerMatcher.inspect(t, w, fsName, plName)
+			headerMatcher.inspect(t, w, r.Context(), fsName, plName)
 			switch {
 			case r.URL.Path == firstRequestPathPanic:
 				close(firstHandlerDoneCh)
@@ -785,7 +785,7 @@ func TestPriorityAndFairnessWithPanicRecoveryAndTimeoutFilter(t *testing.T) {
 		rquestTimesOutPath := "/request/time-out-as-designed"
 		reqHandlerCompletedCh, callerRoundTripDoneCh := make(chan struct{}), make(chan struct{})
 		requestHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			headerMatcher.inspect(t, w, fsName, plName)
+			headerMatcher.inspect(t, w, r.Context(), fsName, plName)
 
 			if r.URL.Path == rquestTimesOutPath {
 				defer close(reqHandlerCompletedCh)
@@ -858,7 +858,7 @@ func TestPriorityAndFairnessWithPanicRecoveryAndTimeoutFilter(t *testing.T) {
 		reqHandlerErrCh, callerRoundTripDoneCh := make(chan error, 1), make(chan struct{})
 		rquestTimesOutPath := "/request/time-out-as-designed"
 		requestHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			headerMatcher.inspect(t, w, fsName, plName)
+			headerMatcher.inspect(t, w, r.Context(), fsName, plName)
 
 			if r.URL.Path == rquestTimesOutPath {
 				<-callerRoundTripDoneCh
@@ -937,7 +937,7 @@ func TestPriorityAndFairnessWithPanicRecoveryAndTimeoutFilter(t *testing.T) {
 		rquestTimesOutPath := "/request/time-out-as-designed"
 		reqHandlerErrCh, callerRoundTripDoneCh := make(chan error, 1), make(chan struct{})
 		requestHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			headerMatcher.inspect(t, w, fsName, plName)
+			headerMatcher.inspect(t, w, r.Context(), fsName, plName)
 
 			if r.URL.Path == rquestTimesOutPath {
 
@@ -1016,7 +1016,7 @@ func TestPriorityAndFairnessWithPanicRecoveryAndTimeoutFilter(t *testing.T) {
 		firstReqHandlerErrCh, firstReqInProgressCh := make(chan error, 1), make(chan struct{})
 		firstReqRoundTripDoneCh, secondReqRoundTripDoneCh := make(chan struct{}), make(chan struct{})
 		requestHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			headerMatcher.inspect(t, w, fsName, plName)
+			headerMatcher.inspect(t, w, r.Context(), fsName, plName)
 			switch {
 			case r.URL.Path == firstRequestTimesOutPath:
 				close(firstReqInProgressCh)
@@ -1188,28 +1188,35 @@ func newHTTP2ServerWithClient(handler http.Handler, clientTimeout time.Duration)
 type headerMatcher struct{}
 
 // verifies that the expected flow schema and priority level UIDs are attached to the header.
-func (m *headerMatcher) inspect(t *testing.T, w http.ResponseWriter, expectedFS, expectedPL string) {
+func (m *headerMatcher) inspect(t *testing.T, w http.ResponseWriter, ctx context.Context, expectedFS, expectedPL string) {
 	t.Helper()
-	err := func() error {
-		if w == nil {
-			return fmt.Errorf("expected a non nil HTTP response")
-		}
 
-		key := flowcontrol.ResponseHeaderMatchedFlowSchemaUID
-		if value := w.Header().Get(key); expectedFS != value {
-			return fmt.Errorf("expected HTTP header %s to have value %q, but got: %q", key, expectedFS, value)
-		}
-
-		key = flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID
-		if value := w.Header().Get(key); expectedPL != value {
-			return fmt.Errorf("expected HTTP header %s to have value %q, but got %q", key, expectedPL, value)
-		}
-		return nil
-	}()
-	if err == nil {
+	if w == nil {
+		t.Errorf("expected a non nil HTTP response")
 		return
 	}
-	t.Errorf("Expected APF headers to match, but got: %v", err)
+
+	checkHeader := func(key, expected string) {
+		actual := w.Header().Get(key)
+		if actual == expected {
+			return
+		}
+
+		// Header writes are best-effort once the request context is canceled.
+		if actual == "" {
+			select {
+			case <-ctx.Done():
+				t.Logf("Skipping APF header assertion for %s: request context already done", key)
+				return
+			default:
+			}
+		}
+
+		t.Errorf("expected HTTP header %s to have value %q, but got: %q", key, expected, actual)
+	}
+
+	checkHeader(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, expectedFS)
+	checkHeader(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, expectedPL)
 }
 
 // when a request panics, http2 resets the stream with an INTERNAL_ERROR message