diff --git a/internal/agentcfg/elasticsearch_test.go b/internal/agentcfg/elasticsearch_test.go index 51432f805e6..4671934c847 100644 --- 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,