move lazy geoipdb download to test setup

Author: rubvs

systemtest/agentconfig_test.go

@@ -39,7 +39,6 @@ import (
 
 func TestAgentConfig(t *testing.T) {
 	systemtest.CleanupElasticsearch(t)
-	systemtest.GeoIpLazyDownload(t)
 
 	serviceName := "systemtest_service"
 	serviceEnvironment := "testing"

systemtest/aggregation_test.go

@@ -39,8 +39,6 @@ import (
 
 func TestTransactionAggregation(t *testing.T) {
 	systemtest.CleanupElasticsearch(t)
-	systemtest.GeoIpLazyDownload(t)
-
 	srv := apmservertest.NewUnstartedServerTB(t)
 	require.NoError(t, srv.Start())
 
@@ -120,8 +118,6 @@ func TestTransactionAggregation(t *testing.T) {
 
 func TestTransactionAggregationShutdown(t *testing.T) {
 	systemtest.CleanupElasticsearch(t)
-	systemtest.GeoIpLazyDownload(t)
-
 	srv := apmservertest.NewServerTB(t)
 
 	// Send a transaction to the server to be aggregated.
@@ -150,8 +146,6 @@ func TestTransactionAggregationShutdown(t *testing.T) {
 func TestServiceDestinationAggregation(t *testing.T) {
 	t.Setenv("ELASTIC_APM_GLOBAL_LABELS", "department_name=apm,organization=observability,company=elastic")
 	systemtest.CleanupElasticsearch(t)
-	systemtest.GeoIpLazyDownload(t)
-
 	srv := apmservertest.NewServerTB(t)
 
 	// Send spans to the server to be aggregated.

systemtest/geoip.go

@@ -18,9 +18,17 @@
 package systemtest
 
 import (
+	"encoding/json"
+	"fmt"
+	"io"
 	"log"
+	"net/http"
 	"os"
+	"strings"
 	"testing"
+	"time"
+
+	"github.com/elastic/apm-server/systemtest/apmservertest"
 )
 
 func TestMain(m *testing.M) {
@@ -30,9 +38,90 @@ func TestMain(m *testing.M) {
 		log.Fatalf("failed to start stack containers: %v", err)
 	}
 	initElasticSearch()
+	if err := GeoIpLazyDownload(); err != nil {
+		log.Fatalf("failed to download geoip database: %v", err)
+	}
 	initKibana()
 	initSettings()
 	initOTEL()
 	log.Println("INFO: running system tests...")
 	os.Exit(m.Run())
 }
+
+var stats struct {
+	GeoIPStats `json:"stats"`
+}
+
+type GeoIPStats struct {
+	DatabasesCount int `json:"databases_count"`
+}
+
+const requestBody = `{"metadata": { "service": {"name": "alice", "agent": {"version": "3.14.0", "name": "elastic-node"}}}}
+{"metricset": {"samples":{"a":{"value":3.2}}, "timestamp": 1496170422281000}}`
+
+func GeoIpLazyDownload() error {
+	srv := apmservertest.NewUnstartedServer()
+	if err := srv.Start(); err != nil {
+		return err
+	}
+
+	serverURL := srv.URL + "/intake/v2/events"
+	req, err := http.NewRequest("POST", serverURL, strings.NewReader(requestBody))
+	if err != nil {
+		return err
+	}
+
+	req.Header.Set("Content-Type", "application/x-ndjson")
+	req.Header.Set("X-Forwarded-For", "8.8.8.8")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	if resp.StatusCode != http.StatusAccepted {
+		return fmt.Errorf("query failed with status code: %d", resp.StatusCode)
+	}
+
+	io.Copy(io.Discard, resp.Body)
+	resp.Body.Close()
+
+	if err := waitGeoIPDownload(); err != nil {
+		return err
+	}
+
+	// Clean all datastreams containing tags.
+	cleanupElasticsearch()
+
+	return nil
+}
+
+func waitGeoIPDownload() error {
+	timer := time.NewTimer(2 * time.Minute)
+	ticker := time.NewTicker(500 * time.Millisecond)
+	defer ticker.Stop()
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			resp, err := Elasticsearch.Ingest.GeoIPStats()
+			if err != nil {
+				return err
+			}
+			body, err := io.ReadAll(resp.Body)
+			resp.Body.Close()
+			if err != nil {
+				return err
+			}
+			if err := json.Unmarshal(body, &stats); err != nil {
+				return err
+			}
+			if stats.DatabasesCount == 3 {
+				log.Println("GeoIP database downloaded")
+				return nil
+			}
+		case <-timer.C:
+			return fmt.Errorf("download timeout exceeded")
+		}
+	}
+}

systemtest/main_test.go

@@ -40,7 +40,6 @@ import (
 
 func TestRUMXForwardedFor(t *testing.T) {
 	systemtest.CleanupElasticsearch(t)
-	systemtest.GeoIpLazyDownload(t)
 
 	srv := apmservertest.NewUnstartedServerTB(t)
 	srv.Config.RUM = &apmservertest.RUMConfig{Enabled: true}