gorz/internal/proxy/loadbalancer_test.go

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)
	}
}