From 26a957b340a9e4bb40859f85c5f3fadf1684e04d Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Wed, 30 Jun 2021 19:45:30 +0800
Subject: [PATCH] sourcemap: fix flaky test (#5581) (#5588)

The test was failing because the dial attempt would
sometimes timeout and return a different error.

Change the test to first wait for the request to be
received by the server, and then check that context
cancellation is honoured. We do this by explicitly
cancelling the context, rather than using a timeout.

(cherry picked from commit 062cb9cc1f0e48e7b9c126bba8ad10671cbc04b9)

Co-authored-by: Andrew Wilkins <axw@elastic.co>
---
 sourcemap/store_test.go | 55 +++++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 18 deletions(-)

diff --git a/sourcemap/store_test.go b/sourcemap/store_test.go
index d29f5baa695..2ee44cd4480 100644
--- a/sourcemap/store_test.go
+++ b/sourcemap/store_test.go
@@ -153,22 +153,24 @@ func TestStore_Fetch(t *testing.T) {
 	})
 }
 
-func TestFetchTimeout(t *testing.T) {
-	// TODO(stn): fix this flaky test
-	t.Skip()
+func TestFetchContext(t *testing.T) {
 	var (
-		errs int64
-
-		apikey      = "supersecret"
-		name        = "webapp"
-		version     = "1.0.0"
-		path        = "/my/path/to/bundle.js.map"
-		c           = http.DefaultClient
-		ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond)
+		apikey  = "supersecret"
+		name    = "webapp"
+		version = "1.0.0"
+		path    = "/my/path/to/bundle.js.map"
+		c       = http.DefaultClient
 	)
-	defer cancel()
+
+	requestReceived := make(chan struct{})
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		<-ctx.Done()
+		select {
+		case requestReceived <- struct{}{}:
+		case <-r.Context().Done():
+			return
+		}
+		// block until the client cancels the request
+		<-r.Context().Done()
 	}))
 	defer ts.Close()
 
@@ -190,13 +192,30 @@ func TestFetchTimeout(t *testing.T) {
 	assert.NoError(t, err)
 	logger := logp.NewLogger(logs.Sourcemap)
 	store, err := newStore(b, logger, time.Minute)
-	assert.NoError(t, err)
+	require.NoError(t, err)
 
-	_, err = store.Fetch(ctx, name, version, path)
-	assert.True(t, errors.Is(err, context.DeadlineExceeded))
-	atomic.AddInt64(&errs, 1)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	fetchReturned := make(chan error, 1)
+	go func() {
+		defer close(fetchReturned)
+		_, err := store.Fetch(ctx, name, version, path)
+		fetchReturned <- err
+	}()
+	select {
+	case <-requestReceived:
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for server to receive request")
+	}
 
-	assert.Equal(t, int64(1), errs)
+	// Check that cancelling the context unblocks the request.
+	cancel()
+	select {
+	case err := <-fetchReturned:
+		assert.True(t, errors.Is(err, context.Canceled))
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for Fetch to return")
+	}
 }
 
 func TestConcurrentFetch(t *testing.T) {