Function runtimeAPIServer

lambda/invoke_loop_test.go:484–543

(eventPayload string, failAfter int, overrides ...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{}

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			numInvokesRequestedLock.Lock()
			metadata := defaultInvokeMetadata()
			if numInvokesRequested < len(overrides) {
				metadata = overrides[numInvokesRequested]
			}
			numInvokesRequested++
			shouldFail := numInvokesRequested > failAfter
			numInvokesRequestedLock.Unlock()
			record.lock.Lock()
			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
			}
			w.Header().Add(string(headerAWSRequestID), metadata.requestID)
			w.Header().Add(string(headerDeadlineMS), metadata.deadline)
			w.Header().Add(string(headerInvokedFunctionARN), metadata.functionARN)
			w.Header().Add(string(headerClientContext), metadata.clientContext)
			w.Header().Add(string(headerCognitoIdentity), metadata.cognito)
			w.Header().Add(string(headerTraceID), metadata.xray)
			if metadata.tenantID != "" {
				w.Header().Add(string(headerTenantID), metadata.tenantID)
			}
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(eventPayload))
		case http.MethodPost:
			response := bytes.NewBuffer(nil)
			_, _ = io.Copy(response, r.Body)
			_ = r.Body.Close()
			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)
		}
	}))