FastDeploy/fastdeploy/golang_router/internal/manager/handler_test.go

package manager

import (
	"context"
	"testing"

	"github.com/PaddlePaddle/FastDeploy/router/internal/config"
	"github.com/stretchr/testify/assert"
)

func TestInit(t *testing.T) {
	cfg := &config.Config{
		Server: config.ServerConfig{
			Splitwise: true,
		},
		Manager: config.ManagerConfig{
			HealthCheckTimeoutSecs: 5.0,
			HealthCheckEndpoint:    "/health",
			HealthFailureThreshold: 3,
			HealthSuccessThreshold: 2,
		},
	}

	Init(cfg)

	assert.NotNil(t, DefaultManager)
	assert.True(t, DefaultManager.splitwise)
	assert.Equal(t, "/health", healthEndpoint)
	assert.Equal(t, 3, failureThreshold)
	assert.Equal(t, 2, successThreshold)
}

func TestWorkerMapToList(t *testing.T) {
	// Setup test data
	Init(&config.Config{})
	DefaultManager.prefillWorkerMap = map[string]*WorkerInfo{
		"http://worker1": {Url: "http://worker1"},
		"http://worker2": {Url: "http://worker2"},
	}
	DefaultManager.decodeWorkerMap = map[string]*WorkerInfo{
		"http://worker3": {Url: "http://worker3"},
	}
	DefaultManager.mixedWorkerMap = map[string]*WorkerInfo{
		"http://worker4": {Url: "http://worker4"},
	}

	t.Run("prefill workers", func(t *testing.T) {
		workers := WorkerMapToList(context.Background(), "prefill")
		assert.Len(t, workers, 2)
		assert.Contains(t, workers, "http://worker1")
		assert.Contains(t, workers, "http://worker2")
	})

	t.Run("decode workers", func(t *testing.T) {
		workers := WorkerMapToList(context.Background(), "decode")
		assert.Len(t, workers, 1)
		assert.Contains(t, workers, "http://worker3")
	})

	t.Run("mixed workers", func(t *testing.T) {
		workers := WorkerMapToList(context.Background(), "mixed")
		assert.Len(t, workers, 1)
		assert.Contains(t, workers, "http://worker4")
	})

	t.Run("invalid worker type", func(t *testing.T) {
		workers := WorkerMapToList(context.Background(), "invalid")
		assert.Len(t, workers, 0)
	})
}

func TestManager_GetHealthyURLs(t *testing.T) {
	// Setup test data
	Init(&config.Config{})
	DefaultManager.prefillWorkerMap = map[string]*WorkerInfo{
		"worker1": {Url: "http://worker1"},
	}
	DefaultManager.decodeWorkerMap = map[string]*WorkerInfo{
		"worker2": {Url: "http://worker2"},
	}
	DefaultManager.mixedWorkerMap = map[string]*WorkerInfo{
		"worker3": {Url: "http://worker3"},
	}

	urls := DefaultManager.GetHealthyURLs(context.Background())
	assert.Len(t, urls, 3)
	assert.Contains(t, urls, "worker1")
	assert.Contains(t, urls, "worker2")
	assert.Contains(t, urls, "worker3")
}

func TestSelectWorker(t *testing.T) {
	// Setup test data
	Init(&config.Config{})
	DefaultManager.mixedWorkerMap = map[string]*WorkerInfo{
		"http://worker1": {Url: "http://worker1"},
	}

	// This will fail because SelectWorker depends on scheduler
	// which we don't want to mock in this unit test
	t.Skip("Integration test requiring scheduler setup")
}

func TestSelectWorkerPair(t *testing.T) {
	// Setup test data
	Init(&config.Config{})
	DefaultManager.prefillWorkerMap = map[string]*WorkerInfo{
		"http://worker1": {Url: "http://worker1"},
	}
	DefaultManager.decodeWorkerMap = map[string]*WorkerInfo{
		"http://worker2": {Url: "http://worker2"},
	}

	// This will fail because SelectWorkerPair depends on scheduler
	// which we don't want to mock in this unit test
	t.Skip("Integration test requiring scheduler setup")
}