Merge pull request #64 from iflytek/feat/statedriven&reconciler

add unit tests for basic components

Author: bestksl

go.mod

@@ -59,6 +59,7 @@ require (
 	github.com/spf13/afero v1.15.0 // indirect
 	github.com/spf13/cast v1.10.0 // indirect
 	github.com/spf13/pflag v1.0.10 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.3.0 // indirect

internal/core/orchestrator/orchestrator.go

@@ -13,11 +13,15 @@ import (
 	"fmt"
 )
 
+// CurrentShimletGetter 定义获取当前shimlet ID的函数类型
+type CurrentShimletGetter func() string
+
 type Orchestrator struct {
-	shimReg    *typereg.TypeReg[shimlet.Shimlet]
-	goalSetReg map[string]*goal.GoalSet
-	specStore  spec.Store
-	queue      *workqueue.Queue
+	shimReg             *typereg.TypeReg[shimlet.Shimlet]
+	goalSetReg          map[string]*goal.GoalSet
+	specStore           spec.Store
+	queue               *workqueue.Queue
+	currentShimletGetter CurrentShimletGetter
 }
 
 func NewOrchestrator(
@@ -27,10 +31,31 @@ func NewOrchestrator(
 	specStore spec.Store,
 ) *Orchestrator {
 	return &Orchestrator{
-		queue:      queue,
-		shimReg:    shimReg,
-		goalSetReg: pipeReg,
-		specStore:  specStore,
+		queue:               queue,
+		shimReg:             shimReg,
+		goalSetReg:          pipeReg,
+		specStore:           specStore,
+		currentShimletGetter: func() string {
+			return config.Get().CurrentShimlet
+		},
+	}
+}
+
+// NewOrchestratorWithShimletGetter 创建一个可以自定义currentShimletGetter的Orchestrator实例
+// 主要用于测试
+func NewOrchestratorWithShimletGetter(
+	shimReg *typereg.TypeReg[shimlet.Shimlet],
+	pipeReg map[string]*goal.GoalSet,
+	queue *workqueue.Queue,
+	specStore spec.Store,
+	currentShimletGetter CurrentShimletGetter,
+) *Orchestrator {
+	return &Orchestrator{
+		queue:               queue,
+		shimReg:             shimReg,
+		goalSetReg:          pipeReg,
+		specStore:           specStore,
+		currentShimletGetter: currentShimletGetter,
 	}
 }
 
@@ -46,7 +71,7 @@ func (o *Orchestrator) Provision(spec *dto.RequirementSpec) error {
 
 	// RequirementSpec 持久化 部署期望
 	spec.ReplicaCount = 1
-	spec.ShimletName = config.Get().CurrentShimlet
+	spec.ShimletName = o.currentShimletGetter()
 	// 如果这里是更新, 则需要 对应goalset reconcile 检测到 不一致 并调用ensure 闭环
 	o.specStore.Set(spec.ServiceId, spec)
 
@@ -59,7 +84,7 @@ func (o *Orchestrator) Provision(spec *dto.RequirementSpec) error {
 // DeleteService 删除指定的模型服务
 func (o *Orchestrator) DeleteService(serviceID string) error {
 	// 获取当前使用的shimlet
-	currentShimletId := config.Get().CurrentShimlet
+	currentShimletId := o.currentShimletGetter()
 	runtimeShimlet, err := o.shimReg.GetSingleton(currentShimletId)
 	if err != nil {
 		log.Error("get runtime shimlet error", err)
@@ -82,7 +107,7 @@ func (o *Orchestrator) GetServiceStatus(serviceID string) (*dto.RuntimeStatus, e
 		return nil, fmt.Errorf("serviceID is required")
 	}
 
-	currentShimletId := config.Get().CurrentShimlet
+	currentShimletId := o.currentShimletGetter()
 	runtimeShimlet, err := o.shimReg.GetSingleton(currentShimletId)
 	if err != nil {
 		return nil, err
@@ -95,4 +120,4 @@ func (o *Orchestrator) GetServiceStatus(serviceID string) (*dto.RuntimeStatus, e
 		status.EndPoint += "/v1/chat/completions"
 	}
 	return status, nil
-}
+}

internal/core/orchestrator/orchestrator_test.go

@@ -1 +1,131 @@
 package orchestrator
+
+import (
+	"astron-xmod-shim/internal/config"
+	"astron-xmod-shim/internal/core/goal"
+	"astron-xmod-shim/internal/core/shimlet"
+	"astron-xmod-shim/internal/core/spec"
+	"astron-xmod-shim/internal/core/typereg"
+	"astron-xmod-shim/internal/core/workqueue"
+	dto "astron-xmod-shim/internal/dto/deploy"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+// MockShimlet 是shimlet.Shimlet接口的简化mock实现
+type MockShimlet struct {
+	mock.Mock
+}
+
+func (m *MockShimlet) ID() string {
+	args := m.Called()
+	return args.String(0)
+}
+
+func (m *MockShimlet) Description() string {
+	args := m.Called()
+	return args.String(0)
+}
+
+func (m *MockShimlet) InitWithConfig(confPath string) error {
+	args := m.Called(confPath)
+	return args.Error(0)
+}
+
+func (m *MockShimlet) Apply(spec *dto.RequirementSpec) error {
+	args := m.Called(spec)
+	return args.Error(0)
+}
+
+func (m *MockShimlet) Delete(resourceId string) error {
+	args := m.Called(resourceId)
+	return args.Error(0)
+}
+
+func (m *MockShimlet) Status(resourceId string) (*dto.RuntimeStatus, error) {
+	args := m.Called(resourceId)
+	status, _ := args.Get(0).(*dto.RuntimeStatus)
+	return status, args.Error(1)
+}
+
+func (m *MockShimlet) ListDeployedServices() ([]string, error) {
+	args := m.Called()
+	services, _ := args.Get(0).([]string)
+	return services, args.Error(1)
+}
+
+// 简化的测试用例
+func TestOrchestrator(t *testing.T) {
+	// 创建临时配置文件
+	configContent := `
+current-shimlet: "k8s"
+shimlets:
+  k8s:
+    config-path: "/conf/k8s-shimlet.yaml"
+`
+	tmpFile, err := os.CreateTemp("", "test-config-*.yaml")
+	assert.NoError(t, err)
+	defer os.Remove(tmpFile.Name())
+
+	_, err = tmpFile.WriteString(configContent)
+	assert.NoError(t, err)
+	assert.NoError(t, tmpFile.Close())
+
+	// 设置配置文件路径
+	config.SetConfigPath(tmpFile.Name())
+
+	// 测试Provision方法
+	t.Run("Provision", func(t *testing.T) {
+		// 创建测试需要的组件
+		mockQueue := workqueue.New()
+		defer mockQueue.ShutDown()
+
+		mockSpecStore := spec.NewMemoryStore()
+		shimReg := typereg.New[shimlet.Shimlet]()
+
+		// 创建并配置mock shimlet
+		mockShim := new(MockShimlet)
+		mockShim.On("ID").Return("k8s")
+		mockShim.On("InitWithConfig", "/conf/k8s-shimlet.yaml").Return(nil)
+
+		// 注册mock shimlet到注册中心
+		shimReg.AutoRegister(mockShim)
+
+		// 创建orchestrator实例，使用固定的shimlet ID "k8s"
+		orchestrator := NewOrchestratorWithShimletGetter(
+			shimReg,
+			map[string]*goal.GoalSet{},
+			mockQueue,
+			mockSpecStore,
+			func() string {
+				return "k8s"
+			},
+		)
+
+		serviceSpec := &dto.RequirementSpec{
+			ServiceId:            "test-service",
+			GoalSetName:          "test-goalset",
+			ResourceRequirements: &dto.ResourceRequirements{},
+		}
+
+		err := orchestrator.Provision(serviceSpec)
+		assert.NoError(t, err)
+
+		// 验证规格被保存
+		savedSpec := mockSpecStore.Get("test-service")
+		assert.NotNil(t, savedSpec)
+
+		// 验证项目被添加到队列
+		assert.Equal(t, 1, mockQueue.Len())
+
+		// 清理队列
+		if mockQueue.Len() > 0 {
+			_, done := mockQueue.Get()
+			done()
+		}
+	})
+
+}

internal/core/spec/memory_store_test.go

@@ -0,0 +1,78 @@
+package spec
+
+import (
+	dto "astron-xmod-shim/internal/dto/deploy"
+	"testing"
+)
+
+func TestMemoryStoreBasicOperations(t *testing.T) {
+	// 创建内存存储
+	store := NewMemoryStore()
+
+	// 创建测试数据
+	serviceID := "test-service"
+	spec := &dto.RequirementSpec{
+		ServiceId:    serviceID,
+		ModelName:    "test-model",
+		ReplicaCount: 1,
+	}
+
+	// 测试设置和获取
+	store.Set(serviceID, spec)
+	retrieved := store.Get(serviceID)
+	if retrieved == nil {
+		t.Error("Expected to retrieve non-nil spec")
+	} else if retrieved.ServiceId != serviceID {
+		t.Errorf("Expected service ID '%s', got '%s'", serviceID, retrieved.ServiceId)
+	} else if retrieved.ModelName != "test-model" {
+		t.Errorf("Expected model name 'test-model', got '%s'", retrieved.ModelName)
+	}
+
+	// 测试删除
+	store.Delete(serviceID)
+	retrieved = store.Get(serviceID)
+	if retrieved != nil {
+		t.Error("Expected nil after delete")
+	}
+}
+
+func TestMemoryStoreMultipleServices(t *testing.T) {
+	// 创建内存存储
+	store := NewMemoryStore()
+
+	// 创建多个服务规格
+	spec1 := &dto.RequirementSpec{ServiceId: "service-1", ModelName: "model-1"}
+	spec2 := &dto.RequirementSpec{ServiceId: "service-2", ModelName: "model-2"}
+	spec3 := &dto.RequirementSpec{ServiceId: "service-3", ModelName: "model-3"}
+
+	// 设置多个服务
+	store.Set("service-1", spec1)
+	store.Set("service-2", spec2)
+	store.Set("service-3", spec3)
+
+	// 验证所有服务都能正确获取
+	if retrieved := store.Get("service-1"); retrieved == nil || retrieved.ModelName != "model-1" {
+		t.Error("Failed to retrieve service-1")
+	}
+	if retrieved := store.Get("service-2"); retrieved == nil || retrieved.ModelName != "model-2" {
+		t.Error("Failed to retrieve service-2")
+	}
+	if retrieved := store.Get("service-3"); retrieved == nil || retrieved.ModelName != "model-3" {
+		t.Error("Failed to retrieve service-3")
+	}
+
+	// 验证不存在的服务返回nil
+	if retrieved := store.Get("non-existent"); retrieved != nil {
+		t.Error("Expected nil for non-existent service")
+	}
+
+	// 删除一个服务并验证
+	store.Delete("service-2")
+	if retrieved := store.Get("service-2"); retrieved != nil {
+		t.Error("Expected nil after deleting service-2")
+	}
+	// 确保其他服务仍然存在
+	if retrieved := store.Get("service-1"); retrieved == nil {
+		t.Error("service-1 should still exist")
+	}
+}