From 8d3e975faf7640ae95ab61f2085cfe6a3903191d Mon Sep 17 00:00:00 2001
From: ICHINOSE Shogo <shogo82148@gmail.com>
Date: Thu, 11 Dec 2025 10:25:57 +0900
Subject: [PATCH 1/3] fix flaky test TestRuntimeAPILoopWithConcurrency

---
 lambda/invoke_loop_test.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go
index 7c1870a8..af0e0c0c 100644
--- a/lambda/invoke_loop_test.go
+++ b/lambda/invoke_loop_test.go
@@ -442,6 +442,7 @@ func defaultInvokeMetadata() eventMetadata {
 }
 
 func runtimeAPIServer(eventPayload string, failAfter int, overrides ...eventMetadata) (*httptest.Server, *requestRecord) {
+	ctx, cancel := context.WithCancel(context.Background())
 	numInvokesRequested := 0
 	numInvokesRequestedLock := sync.Mutex{}
 	record := &requestRecord{}
@@ -461,6 +462,7 @@ func runtimeAPIServer(eventPayload string, failAfter int, overrides ...eventMeta
 			record.nGets++
 			record.lock.Unlock()
 			if shouldFail {
+				<-ctx.Done() // wait for all handlers to finish.
 				w.WriteHeader(http.StatusGone)
 				_, _ = w.Write([]byte("END THE TEST!"))
 				return
@@ -483,10 +485,15 @@ func runtimeAPIServer(eventPayload string, failAfter int, overrides ...eventMeta
 			w.WriteHeader(http.StatusAccepted)
 			record.lock.Lock()
 			record.nPosts++
+			done := record.nPosts >= failAfter
 			record.responses = append(record.responses, response.Bytes())
 			record.contentTypes = append(record.contentTypes, r.Header.Get("Content-Type"))
 			record.xrayCauses = append(record.xrayCauses, r.Header.Get(headerXRayErrorCause))
 			record.lock.Unlock()
+			if done {
+				// all handlers are done, cancel the context to let the GET handler exit.
+				cancel()
+			}
 		default:
 			w.WriteHeader(http.StatusBadRequest)
 		}

From 137df52d3a6bdf9a690720f2febd995574c88fe7 Mon Sep 17 00:00:00 2001
From: ICHINOSE Shogo <shogo82148@gmail.com>
Date: Thu, 11 Dec 2025 11:09:48 +0900
Subject: [PATCH 2/3] Add context cancellation handling in runtimeAPIServer
 test

---
 lambda/invoke_loop_test.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go
index af0e0c0c..47af4aea 100644
--- a/lambda/invoke_loop_test.go
+++ b/lambda/invoke_loop_test.go
@@ -448,6 +448,14 @@ func runtimeAPIServer(eventPayload string, failAfter int, overrides ...eventMeta
 	record := &requestRecord{}
 
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-ctx.Done():
+			w.WriteHeader(http.StatusGone)
+			_, _ = w.Write([]byte("END THE TEST!"))
+			return
+		default:
+		}
+
 		switch r.Method {
 		case http.MethodGet:
 			numInvokesRequestedLock.Lock()

From d3e60222007765c65542f6626f1bac3b26a24cee Mon Sep 17 00:00:00 2001
From: ICHINOSE Shogo <shogo82148@gmail.com>
Date: Thu, 11 Dec 2025 11:29:26 +0900
Subject: [PATCH 3/3] Fix concurrency handling in
 TestRuntimeAPILoopWithConcurrency

---
 lambda/invoke_loop_gte_go122_test.go | 2 ++
 lambda/invoke_loop_test.go           | 8 --------
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/lambda/invoke_loop_gte_go122_test.go b/lambda/invoke_loop_gte_go122_test.go
index cf3a906a..83284562 100644
--- a/lambda/invoke_loop_gte_go122_test.go
+++ b/lambda/invoke_loop_gte_go122_test.go
@@ -65,8 +65,10 @@ func TestRuntimeAPILoopWithConcurrency(t *testing.T) {
 	endpoint := strings.Split(ts.URL, "://")[1]
 	expectedError := fmt.Sprintf("failed to GET http://%s/2018-06-01/runtime/invocation/next: got unexpected status code: 410", endpoint)
 	assert.EqualError(t, startRuntimeAPILoopWithConcurrency(endpoint, handler, concurrency), expectedError)
+	record.lock.Lock()
 	assert.GreaterOrEqual(t, record.nGets, nInvokes+1)
 	assert.Equal(t, nInvokes, record.nPosts)
+	record.lock.Unlock()
 	assert.Equal(t, int32(concurrency), maxActive.Load())
 	responses := make(map[string]int)
 	for _, response := range record.responses {
diff --git a/lambda/invoke_loop_test.go b/lambda/invoke_loop_test.go
index 47af4aea..af0e0c0c 100644
--- a/lambda/invoke_loop_test.go
+++ b/lambda/invoke_loop_test.go
@@ -448,14 +448,6 @@ func runtimeAPIServer(eventPayload string, failAfter int, overrides ...eventMeta
 	record := &requestRecord{}
 
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		select {
-		case <-ctx.Done():
-			w.WriteHeader(http.StatusGone)
-			_, _ = w.Write([]byte("END THE TEST!"))
-			return
-		default:
-		}
-
 		switch r.Method {
 		case http.MethodGet:
 			numInvokesRequestedLock.Lock()