fix flaky test TestRuntimeAPILoopWithConcurrency (#606)

Author: shogo82148

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 {

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)
 		}