functionaltests/internal/asserts/doc_count.go

// Licensed to Elasticsearch B.V. under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Elasticsearch B.V. licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package asserts

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/elastic/apm-server/functionaltests/internal/esclient"
)

// CheckDocCount checks if difference in data stream document count between
// current and previous state is equal to the expected, barring the skipped
// data streams.
//
// NOTE: Ingestion should never remove documents. If the expectedDiff is
// negative, the data stream will be expected to appear, but the document
// count will not be asserted.
func CheckDocCount(
	t *testing.T,
	currDocCount, prevDocCount, expectedDiff esclient.DataStreamsDocCount,
) {
	t.Helper()

	if prevDocCount == nil {
		prevDocCount = map[string]int{}
	}

	// Check that all expected data streams appear.
	for ds := range expectedDiff {
		if _, ok := currDocCount[ds]; !ok {
			t.Errorf("expected data stream %s not found", ds)
			continue
		}
	}

	// Check document counts for all data streams.
	for ds, v := range currDocCount {
		e, ok := expectedDiff[ds]
		if !ok {
			t.Errorf("unexpected documents (%d) for data stream %s", v, ds)
			continue
		}

		if e < 0 {
			continue
		}

		assert.Equal(t, e, v-prevDocCount[ds],
			fmt.Sprintf("wrong document count difference for data stream %s", ds))
	}
}