Test agentcfg es fetcher scroll id (#13959)

Follow up on #13958 to add test
elastic · Sep 3, 2024 · 2aec4d1 · 2aec4d1
1 parent ecffa8e
commit 2aec4d1
Showing 1 changed file with 21 additions and 8 deletions.
diff --git a/internal/agentcfg/elasticsearch_test.go b/internal/agentcfg/elasticsearch_test.go
@@ -20,7 +20,6 @@ package agentcfg
 import (
 	"context"
 	"encoding/json"
-	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -39,11 +38,10 @@ var sampleHits = []map[string]interface{}{
 	{"_id": "hvKmzYQBfJ4l0GgqXgJt", "_index": ".apm-agent-configuration", "_score": 1, "_source": map[string]interface{}{"@timestamp": 1.669897543277e+12, "applied_by_agent": false, "etag": "2da2f86251165ccced5c5e41100a216b0c880db4", "service": map[string]interface{}{"name": "second"}, "settings": map[string]interface{}{"sanitize_field_names": "foo,bar,baz", "transaction_sample_rate": "0.1"}}},
 }
 
-func newMockElasticsearchClient(t testing.TB, statusCode int, responseFunc func(io.Writer)) *elasticsearch.Client {
+func newMockElasticsearchClient(t testing.TB, handler func(http.ResponseWriter, *http.Request)) *elasticsearch.Client {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("X-Elastic-Product", "Elasticsearch")
-		w.WriteHeader(statusCode)
-		responseFunc(w)
+		handler(w, r)
 	}))
 	t.Cleanup(srv.Close)
 	config := elasticsearch.DefaultConfig()
@@ -80,7 +78,15 @@ func newElasticsearchFetcher(
 
 	i := 0
 
-	fetcher := NewElasticsearchFetcher(newMockElasticsearchClient(t, 200, func(w io.Writer) {
+	fetcher := NewElasticsearchFetcher(newMockElasticsearchClient(t, func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/_search/scroll":
+			scrollID := r.URL.Query().Get("scroll_id")
+			assert.Equal(t, respTmpl["_scroll_id"], scrollID)
+		case "/.apm-agent-configuration/_search":
+		default:
+			assert.Failf(t, "unexpected path", "path: %s", r.URL.Path)
+		}
 		if i < len(hits) {
 			respTmpl["hits"].(map[string]interface{})["hits"] = hits[i : i+searchSize]
 		} else {
@@ -89,6 +95,7 @@ func newElasticsearchFetcher(
 
 		b, err := json.Marshal(respTmpl)
 		require.NoError(t, err)
+		w.WriteHeader(200)
 		w.Write(b)
 		i += searchSize
 	}), time.Second, nil, rt)
@@ -173,7 +180,9 @@ func TestFetchUseFallback(t *testing.T) {
 		return Result{}, nil
 	})
 	fetcher := NewElasticsearchFetcher(
-		newMockElasticsearchClient(t, 404, func(w io.Writer) {}),
+		newMockElasticsearchClient(t, func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(404)
+		}),
 		time.Second,
 		fallbackFetcher,
 		apmtest.NewRecordingTracer().Tracer,
@@ -186,7 +195,9 @@ func TestFetchUseFallback(t *testing.T) {
 
 func TestFetchNoFallbackInvalidESCfg(t *testing.T) {
 	fetcher := NewElasticsearchFetcher(
-		newMockElasticsearchClient(t, 401, func(w io.Writer) {}),
+		newMockElasticsearchClient(t, func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(401)
+		}),
 		time.Second,
 		nil,
 		apmtest.NewRecordingTracer().Tracer,
@@ -200,7 +211,9 @@ func TestFetchNoFallbackInvalidESCfg(t *testing.T) {
 
 func TestFetchNoFallback(t *testing.T) {
 	fetcher := NewElasticsearchFetcher(
-		newMockElasticsearchClient(t, 500, func(w io.Writer) {}),
+		newMockElasticsearchClient(t, func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(500)
+		}),
 		time.Second,
 		nil,
 		apmtest.NewRecordingTracer().Tracer,