allow overriding the expected doc count per test

functionaltests/8_15_test.go

@@ -27,6 +27,7 @@ func TestUpgrade_8_15_4_to_8_16_0(t *testing.T) {
 	runESSUpgradeTest(t, essUpgradeTestCase{
 		from: "8.15.4",
 		to:   "8.16.0",
+
 		beforeUpgradeAfterIngest: checkDatastreamWant{
 			Quantity:     8,
 			PreferIlm:    false,

functionaltests/main_test.go

@@ -47,13 +47,20 @@ const (
 // single run of ingest().
 // Only non aggregation data streams are included, as aggregation ones differs on different
 // runs.
-func expectedIngestForASingleRun() esclient.APMDataStreamsDocCount {
-	return map[string]int{
-		"traces-apm-default":                     15013,
-		"metrics-apm.app.opbeans_python-default": 1437,
-		"metrics-apm.internal-default":           1351,
-		"logs-apm.error-default":                 364,
+func expectedIngestForASingleRun(namespace string) esclient.APMDataStreamsDocCount {
+	res := esclient.APMDataStreamsDocCount{}
+	for k, v := range map[string]int{
+		// Add here data stream names without namespace, it will be added automatically
+		// in the next lines.
+		"traces-apm":                     15013,
+		"metrics-apm.app.opbeans_python": 1437,
+		"metrics-apm.internal":           1351,
+		"logs-apm.error":                 364,
+	} {
+		k := fmt.Sprintf("%s-%s", k, namespace)
+		res[k] = v
 	}
+	return res
 }
 
 // getDocsCountPerDS retrieves document count.

functionaltests/to_8_17_test.go

@@ -33,6 +33,7 @@ func TestUpgradeTo8170_plain(t *testing.T) {
 	runESSUpgradeTest(t, essUpgradeTestCase{
 		from: "8.16.1",
 		to:   "8.17.0",
+
 		beforeUpgradeAfterIngest: checkDatastreamWant{
 			Quantity:     8,
 			PreferIlm:    false,
@@ -74,6 +75,8 @@ func TestUpgradeTo8170_reroute(t *testing.T) {
 		from: "8.16.1",
 		to:   "8.17.0",
 
+		expectedDsDocCountForAIngestRun: expectedIngestForASingleRun("rerouted"),
+
 		setupFn: func(ecc *esclient.Client, _ *kbclient.Client, _ esclient.Config) error {
 			return createRerouteIngestPipelines(t, context.Background(), ecc)
 		},
@@ -116,6 +119,7 @@ func TestUpgradeTo8170_withAPMIntegration(t *testing.T) {
 	runESSUpgradeTest(t, essUpgradeTestCase{
 		from: "8.14.3",
 		to:   "8.17.0",
+
 		// APM integration is always enabled through the Elastic Agent Cloud Policy,
 		// no further setup is necessary.
 		afterUpgrade: func(_ *esclient.Client, kbc *kbclient.Client, _ esclient.Config) error {

functionaltests/upgrade_test.go

@@ -41,6 +41,12 @@ type essUpgradeTestCase struct {
 	from string
 	to   string
 
+	// expectedDsDocCountForAIngestRunis the expect document count for each
+	// data stream that the test should observe after a single ingestion run
+	// of the test data. APM data stream don't usually change but we want to
+	// run tests where we use, for example, a dedicated namespace.
+	expectedDsDocCountForAIngestRun esclient.APMDataStreamsDocCount
+
 	// setupFn allows to specify custom logic to happen after test cluster
 	// boostrap and before any ingestion happens.
 	setupFn additionalFn
@@ -59,6 +65,12 @@ func runESSUpgradeTest(t *testing.T, tc essUpgradeTestCase) {
 	start := time.Now()
 	ctx := context.Background()
 
+	// Initialize the empty value to a sane default to avoid repeating
+	// the default value in each test.
+	if len(tc.expectedDsDocCountForAIngestRun) == 0 {
+		tc.expectedDsDocCountForAIngestRun = expectedIngestForASingleRun("default")
+	}
+
 	t.Log("creating deploment with terraform")
 	tf, err := terraform.New(t, t.Name())
 	require.NoError(t, err)
@@ -121,7 +133,7 @@ func runESSUpgradeTest(t *testing.T, tc essUpgradeTestCase) {
 
 	beforeUpgradeCount, err := getDocsCountPerDS(t, ctx, ecc)
 	require.NoError(t, err)
-	assertDocCount(t, beforeUpgradeCount, previous, expectedIngestForASingleRun())
+	assertDocCount(t, beforeUpgradeCount, previous, tc.expectedDsDocCountForAIngestRun)
 
 	t.Log("check data streams")
 	var dss []types.DataStream
@@ -160,7 +172,7 @@ func runESSUpgradeTest(t *testing.T, tc essUpgradeTestCase) {
 	t.Log("check number of documents")
 	afterUpgradeIngestionCount, err := getDocsCountPerDS(t, ctx, ecc)
 	require.NoError(t, err)
-	assertDocCount(t, afterUpgradeIngestionCount, afterUpgradeCount, expectedIngestForASingleRun())
+	assertDocCount(t, afterUpgradeIngestionCount, afterUpgradeCount, tc.expectedDsDocCountForAIngestRun)
 
 	t.Log("check data streams and verify lazy rollover happened")
 	dss2, err := ecc.GetDataStream(ctx, "*apm*")