From 70976ce1ca2317ed9dfcc65d1e66ff6d33be3ef8 Mon Sep 17 00:00:00 2001
From: kruskal <99559985+kruskall@users.noreply.github.com>
Date: Wed, 31 Aug 2022 17:43:21 +0200
Subject: [PATCH] test: do not send header before processing data

Do not write status header before processing data or it could lead to
flaky tests due to a race when reading apmServerInternals.Data.
Add proper logging and return early if an error was returned.
---
 main_test.go | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/main_test.go b/main_test.go
index 3786a2f2..28cfaa57 100644
--- a/main_test.go
+++ b/main_test.go
@@ -98,8 +98,11 @@ func newMockApmServer(t *testing.T, l *zap.SugaredLogger) (*MockServerInternals,
 	apmServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		decompressedBytes, err := e2eTesting.GetDecompressedBytesFromRequest(r)
 		if err != nil {
+			l.Debugf("failed to read decompressedBytes: %v", err)
 			w.WriteHeader(http.StatusInternalServerError)
+			return
 		}
+
 		l.Debugf("Event type received by mock APM server : %s", string(decompressedBytes))
 		switch APMServerBehavior(decompressedBytes) {
 		case TimelyResponse:
@@ -119,12 +122,12 @@ func newMockApmServer(t *testing.T, l *zap.SugaredLogger) (*MockServerInternals,
 			panic("Server crashed")
 		default:
 		}
+
 		if r.RequestURI == "/intake/v2/events" {
-			w.WriteHeader(http.StatusAccepted)
 			apmServerInternals.Data += string(decompressedBytes)
 			l.Debug("APM Payload processed")
+			w.WriteHeader(http.StatusAccepted)
 		} else if r.RequestURI == "/" {
-			w.WriteHeader(http.StatusOK)
 			infoPayload, err := json.Marshal(ApmInfo{
 				BuildDate:    time.Now(),
 				BuildSHA:     "7814d524d3602e70b703539c57568cba6964fc20",
@@ -132,11 +135,14 @@ func newMockApmServer(t *testing.T, l *zap.SugaredLogger) (*MockServerInternals,
 				Version:      "8.1.2",
 			})
 			if err != nil {
+				l.Debugf("failed to marshal payload: %v", err)
 				w.WriteHeader(http.StatusInternalServerError)
+				return
 			}
-			_, err = w.Write(infoPayload)
-			if err != nil {
-				w.WriteHeader(http.StatusInternalServerError)
+
+			if _, err = w.Write(infoPayload); err != nil {
+				l.Debugf("failed to write payload: %v", err)
+				return
 			}
 		}
 	}))