Add some common allocation strategy tests

cmd/otel-allocator/allocation/allocator_test.go

@@ -0,0 +1,148 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed 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 allocation
+
+import (
+	"testing"
+
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/open-telemetry/opentelemetry-operator/cmd/otel-allocator/target"
+)
+
+func TestSetCollectors(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+		cols := MakeNCollectors(3, 0)
+		allocator.SetCollectors(cols)
+
+		expectedColLen := len(cols)
+		collectors := allocator.Collectors()
+		assert.Len(t, collectors, expectedColLen)
+
+		for _, i := range cols {
+			assert.NotNil(t, collectors[i.Name])
+		}
+	})
+}
+
+func TestSetTargets(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+		targets := MakeNNewTargetsWithEmptyCollectors(3, 0)
+		allocator.SetTargets(targets)
+
+		expectedTargetLen := len(targets)
+		actualTargets := allocator.TargetItems()
+		assert.Len(t, actualTargets, expectedTargetLen)
+	})
+}
+
+func TestCanSetSingleTarget(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+		cols := MakeNCollectors(3, 0)
+		targets := MakeNNewTargetsWithEmptyCollectors(1, 3)
+		allocator.SetCollectors(cols)
+		allocator.SetTargets(targets)
+		actualTargetItems := allocator.TargetItems()
+		assert.Len(t, actualTargetItems, 1)
+		for _, item := range actualTargetItems {
+			assert.NotEmpty(t, item.CollectorName)
+		}
+	})
+}
+
+func TestCanSetTargetsBeforeCollectors(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+		cols := MakeNCollectors(3, 0)
+		targets := MakeNNewTargetsWithEmptyCollectors(1, 3)
+		allocator.SetTargets(targets)
+		allocator.SetCollectors(cols)
+		actualTargetItems := allocator.TargetItems()
+		assert.Len(t, actualTargetItems, 1)
+		for _, item := range actualTargetItems {
+			assert.NotEmpty(t, item.CollectorName)
+		}
+	})
+}
+
+func TestAddingAndRemovingTargets(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+		cols := MakeNCollectors(3, 0)
+		allocator.SetCollectors(cols)
+
+		initTargets := MakeNNewTargets(6, 3, 0)
+
+		// test that targets and collectors are added properly
+		allocator.SetTargets(initTargets)
+
+		// verify
+		expectedTargetLen := len(initTargets)
+		assert.Len(t, allocator.TargetItems(), expectedTargetLen)
+
+		// prepare second round of targets
+		tar := MakeNNewTargets(4, 3, 0)
+
+		// test that fewer targets are found - removed
+		allocator.SetTargets(tar)
+
+		// verify
+		targetItems := allocator.TargetItems()
+		expectedNewTargetLen := len(tar)
+		assert.Len(t, targetItems, expectedNewTargetLen)
+
+		// verify results map
+		for _, i := range tar {
+			_, ok := targetItems[i.Hash()]
+			assert.True(t, ok)
+		}
+	})
+
+}
+
+// Tests that two targets with the same target url and job name but different label set are both added.
+func TestAllocationCollision(t *testing.T) {
+	RunForAllStrategies(t, func(t *testing.T, allocator Allocator) {
+
+		cols := MakeNCollectors(3, 0)
+		allocator.SetCollectors(cols)
+		firstLabels := model.LabelSet{
+			"test": "test1",
+		}
+		secondLabels := model.LabelSet{
+			"test": "test2",
+		}
+		firstTarget := target.NewItem("sample-name", "0.0.0.0:8000", firstLabels, "")
+		secondTarget := target.NewItem("sample-name", "0.0.0.0:8000", secondLabels, "")
+
+		targetList := map[string]*target.Item{
+			firstTarget.Hash():  firstTarget,
+			secondTarget.Hash(): secondTarget,
+		}
+
+		// test that targets and collectors are added properly
+		allocator.SetTargets(targetList)
+
+		// verify
+		targetItems := allocator.TargetItems()
+		expectedTargetLen := len(targetList)
+		assert.Len(t, targetItems, expectedTargetLen)
+
+		// verify results map
+		for _, i := range targetList {
+			_, ok := targetItems[i.Hash()]
+			assert.True(t, ok)
+		}
+	})
+}

cmd/otel-allocator/allocation/consistent_hashing_test.go

@@ -20,18 +20,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestCanSetSingleTarget(t *testing.T) {
-	cols := MakeNCollectors(3, 0)
-	c, _ := New("consistent-hashing", logger)
-	c.SetCollectors(cols)
-	c.SetTargets(MakeNNewTargets(1, 3, 0))
-	actualTargetItems := c.TargetItems()
-	assert.Len(t, actualTargetItems, 1)
-	for _, item := range actualTargetItems {
-		assert.Equal(t, "collector-0", item.CollectorName)
-	}
-}
-
 func TestRelativelyEvenDistribution(t *testing.T) {
 	numCols := 15
 	numItems := 10000

cmd/otel-allocator/allocation/least_weighted_test.go

@@ -20,100 +20,12 @@ import (
 	"math/rand"
 	"testing"
 
-	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/assert"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
-
-	"github.com/open-telemetry/opentelemetry-operator/cmd/otel-allocator/target"
 )
 
 var logger = logf.Log.WithName("unit-tests")
 
-func TestSetCollectors(t *testing.T) {
-	s, _ := New("least-weighted", logger)
-
-	cols := MakeNCollectors(3, 0)
-	s.SetCollectors(cols)
-
-	expectedColLen := len(cols)
-	collectors := s.Collectors()
-	assert.Len(t, collectors, expectedColLen)
-
-	for _, i := range cols {
-		assert.NotNil(t, collectors[i.Name])
-	}
-}
-
-func TestAddingAndRemovingTargets(t *testing.T) {
-	// prepare allocator with initial targets and collectors
-	s, _ := New("least-weighted", logger)
-
-	cols := MakeNCollectors(3, 0)
-	s.SetCollectors(cols)
-
-	initTargets := MakeNNewTargets(6, 3, 0)
-
-	// test that targets and collectors are added properly
-	s.SetTargets(initTargets)
-
-	// verify
-	expectedTargetLen := len(initTargets)
-	assert.Len(t, s.TargetItems(), expectedTargetLen)
-
-	// prepare second round of targets
-	tar := MakeNNewTargets(4, 3, 0)
-
-	// test that fewer targets are found - removed
-	s.SetTargets(tar)
-
-	// verify
-	targetItems := s.TargetItems()
-	expectedNewTargetLen := len(tar)
-	assert.Len(t, targetItems, expectedNewTargetLen)
-
-	// verify results map
-	for _, i := range tar {
-		_, ok := targetItems[i.Hash()]
-		assert.True(t, ok)
-	}
-}
-
-// Tests that two targets with the same target url and job name but different label set are both added.
-func TestAllocationCollision(t *testing.T) {
-	// prepare allocator with initial targets and collectors
-	s, _ := New("least-weighted", logger)
-
-	cols := MakeNCollectors(3, 0)
-	s.SetCollectors(cols)
-	firstLabels := model.LabelSet{
-		"test": "test1",
-	}
-	secondLabels := model.LabelSet{
-		"test": "test2",
-	}
-	firstTarget := target.NewItem("sample-name", "0.0.0.0:8000", firstLabels, "")
-	secondTarget := target.NewItem("sample-name", "0.0.0.0:8000", secondLabels, "")
-
-	targetList := map[string]*target.Item{
-		firstTarget.Hash():  firstTarget,
-		secondTarget.Hash(): secondTarget,
-	}
-
-	// test that targets and collectors are added properly
-	s.SetTargets(targetList)
-
-	// verify
-	targetItems := s.TargetItems()
-	expectedTargetLen := len(targetList)
-	assert.Len(t, targetItems, expectedTargetLen)
-
-	// verify results map
-	for _, i := range targetList {
-		_, ok := targetItems[i.Hash()]
-		assert.True(t, ok)
-	}
-}
-
 func TestNoCollectorReassignment(t *testing.T) {
 	s, _ := New("least-weighted", logger)
 

cmd/otel-allocator/allocation/allocatortest.go → cmd/otel-allocator/allocation/testutils.go

@@ -12,13 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Note: These utilities are used by other packages, which is why they're defined in a non-test file.
+
 package allocation
 
 import (
 	"fmt"
 	"strconv"
+	"testing"
 
 	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
+	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/open-telemetry/opentelemetry-operator/cmd/otel-allocator/target"
 )
@@ -70,3 +75,15 @@ func MakeNNewTargetsWithEmptyCollectors(n int, startingIndex int) map[string]*ta
 	}
 	return toReturn
 }
+
+func RunForAllStrategies(t *testing.T, f func(t *testing.T, allocator Allocator)) {
+	allocatorNames := GetRegisteredAllocatorNames()
+	logger := logf.Log.WithName("unit-tests")
+	for _, allocatorName := range allocatorNames {
+		t.Run(allocatorName, func(t *testing.T) {
+			allocator, err := New(allocatorName, logger)
+			require.NoError(t, err)
+			f(t, allocator)
+		})
+	}
+}