LFG

internal/proxy/loadbalancer_test.go

@@ -0,0 +1,270 @@
+package proxy
+
+import (
+	"testing"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+)
+
+func TestRoundRobinLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Test that targets are selected in round-robin order
+	for i := 0; i < 6; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		expectedTarget := targets[i%3]
+		if target.Name != expectedTarget.Name {
+			t.Errorf("Expected target %s, got %s", expectedTarget.Name, target.Name)
+		}
+	}
+}
+
+func TestRoundRobinLoadBalancerWithUnhealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Test that only healthy targets are selected
+	for i := 0; i < 4; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if target.Name == "target2" {
+			t.Errorf("Selected unhealthy target: %s", target.Name)
+		}
+
+		// Should alternate between target1 and target3
+		if i%2 == 0 && target.Name != "target1" {
+			t.Errorf("Expected target1, got %s", target.Name)
+		}
+		if i%2 == 1 && target.Name != "target3" {
+			t.Errorf("Expected target3, got %s", target.Name)
+		}
+	}
+}
+
+func TestRoundRobinLoadBalancerNoHealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	_, err := lb.NextTarget()
+	if err != ErrNoHealthyTargets {
+		t.Errorf("Expected ErrNoHealthyTargets, got %v", err)
+	}
+}
+
+func TestRoundRobinLoadBalancerUpdateTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRoundRobinLoadBalancer(targets)
+
+	// Get first target
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Update targets
+	newTargets := []config.TargetConfig{
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+	lb.UpdateTargets(newTargets)
+
+	// Test that new targets are selected
+	for i := 0; i < 4; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		expectedTarget := newTargets[i%2]
+		if target.Name != expectedTarget.Name {
+			t.Errorf("Expected target %s, got %s", expectedTarget.Name, target.Name)
+		}
+	}
+}
+
+func TestRandomLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRandomLoadBalancer(targets)
+
+	// Test that targets are selected randomly
+	// We'll just check that we don't get errors and that all targets are eventually selected
+	selected := make(map[string]bool)
+	for i := 0; i < 100; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+		selected[target.Name] = true
+	}
+
+	// Check that all targets were selected at least once
+	for _, target := range targets {
+		if !selected[target.Name] {
+			t.Errorf("Target %s was never selected", target.Name)
+		}
+	}
+}
+
+func TestRandomLoadBalancerWithUnhealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewRandomLoadBalancer(targets)
+
+	// Test that only healthy targets are selected
+	for i := 0; i < 100; i++ {
+		target, err := lb.NextTarget()
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if target.Name == "target2" {
+			t.Errorf("Selected unhealthy target: %s", target.Name)
+		}
+	}
+}
+
+func TestLeastConnectionsLoadBalancer(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	// Initially, both targets should have 0 connections
+	// The first target should be selected
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Now target1 should have 1 connection, target2 should have 0
+	// So target2 should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target2" {
+		t.Errorf("Expected target2, got %s", target.Name)
+	}
+
+	// Now both targets should have 1 connection
+	// target1 should be selected again
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Release a connection from target1
+	lb.(*LeastConnectionsLoadBalancer).ReleaseConnection("target1")
+
+	// Now target1 should have 1 connection, target2 should have 1 connection
+	// But since we released from target1, it should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+}
+
+func TestLeastConnectionsLoadBalancerUpdateTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: true},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	// Get first target
+	target, err := lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target1" {
+		t.Errorf("Expected target1, got %s", target.Name)
+	}
+
+	// Update targets
+	newTargets := []config.TargetConfig{
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: true},
+		{Name: "target3", Address: "http://localhost:8083", Protocol: "http", Weight: 1, Healthy: true},
+	}
+	lb.UpdateTargets(newTargets)
+
+	// Test that new targets are selected
+	// The first new target should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target2" {
+		t.Errorf("Expected target2, got %s", target.Name)
+	}
+
+	// The second new target should be selected
+	target, err = lb.NextTarget()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if target.Name != "target3" {
+		t.Errorf("Expected target3, got %s", target.Name)
+	}
+}
+
+func TestLeastConnectionsLoadBalancerNoHealthyTargets(t *testing.T) {
+	targets := []config.TargetConfig{
+		{Name: "target1", Address: "http://localhost:8081", Protocol: "http", Weight: 1, Healthy: false},
+		{Name: "target2", Address: "http://localhost:8082", Protocol: "http", Weight: 1, Healthy: false},
+	}
+
+	lb := NewLeastConnectionsLoadBalancer(targets)
+
+	_, err := lb.NextTarget()
+	if err != ErrNoHealthyTargets {
+		t.Errorf("Expected ErrNoHealthyTargets, got %v", err)
+	}
+}