package service import ( "context" "sync" "testing" "time" "github.com/1024XEngineer/xinfra/server/internal/model" ) // newTestDeliveryService 创建一个用于测试的 DeliveryService,不需要真实的 DB 和配置。 // 仅适用于测试 pub/sub、缓存等内存逻辑。 func newTestDeliveryService() *DeliveryService { return &DeliveryService{ streams: make(map[string]map[chan DeliveryTaskSnapshot]struct{}), stdoutCache: make(map[string]*stdoutCacheItem), } } func TestSubscribeTask(t *testing.T) { svc := newTestDeliveryService() ctx := context.Background() _ = ctx taskID := "test-task-1" // 订阅任务 ch, cancel := svc.SubscribeTask(taskID) defer cancel() // 验证 channel 已注册 svc.streamMu.Lock() if _, ok := svc.streams[taskID]; !ok { svc.streamMu.Unlock() t.Fatal("SubscribeTask did not register channel in streams map") } svc.streamMu.Unlock() // 模拟 broadcastTask 推送 snapshot snapshot := DeliveryTaskSnapshot{ Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning}, } svc.streamMu.Lock() for ch := range svc.streams[taskID] { select { case ch <- snapshot: default: } } svc.streamMu.Unlock() // 接收推送 select { case received := <-ch: if received.Task == nil || received.Task.ID != taskID { t.Fatalf("received snapshot task ID = %v, want %q", received.Task, taskID) } if received.Task.Status != model.TaskRunning { t.Fatalf("received snapshot status = %q, want %q", received.Task.Status, model.TaskRunning) } case <-time.After(time.Second): t.Fatal("timeout waiting for snapshot from SubscribeTask") } } func TestSubscribeTask_Cancel(t *testing.T) { svc := newTestDeliveryService() taskID := "test-task-cancel" ch, cancel := svc.SubscribeTask(taskID) // 调用 cancel cancel() // 验证 channel 已关闭 select { case _, ok := <-ch: if ok { t.Fatal("channel should be closed after cancel, but got a value") } case <-time.After(time.Second): t.Fatal("timeout waiting for channel close") } // 验证已从 streams 中移除 svc.streamMu.Lock() if subs := svc.streams[taskID]; subs != nil && len(subs) > 0 { svc.streamMu.Unlock() t.Fatal("cancel did not remove channel from streams map") } svc.streamMu.Unlock() } func TestSubscribeTask_MultipleSubscribers(t *testing.T) { svc := newTestDeliveryService() taskID := "test-task-multi" ch1, cancel1 := svc.SubscribeTask(taskID) defer cancel1() ch2, cancel2 := svc.SubscribeTask(taskID) defer cancel2() ch3, cancel3 := svc.SubscribeTask(taskID) defer cancel3() // 验证三个订阅者都已注册 svc.streamMu.Lock() subs := svc.streams[taskID] if len(subs) != 3 { svc.streamMu.Unlock() t.Fatalf("expected 3 subscribers, got %d", len(subs)) } svc.streamMu.Unlock() // 推送 snapshot snapshot := DeliveryTaskSnapshot{ Task: &model.DeliveryTask{ID: taskID, Status: model.TaskFinished}, } svc.streamMu.Lock() for ch := range svc.streams[taskID] { select { case ch <- snapshot: default: } } svc.streamMu.Unlock() // 验证三个订阅者都收到 for i, ch := range []<-chan DeliveryTaskSnapshot{ch1, ch2, ch3} { select { case received := <-ch: if received.Task == nil || received.Task.Status != model.TaskFinished { t.Fatalf("subscriber %d: expected TaskFinished, got %+v", i+1, received) } case <-time.After(time.Second): t.Fatalf("subscriber %d: timeout waiting for snapshot", i+1) } } } func TestSubscribeTask_CancelOneDoesNotAffectOthers(t *testing.T) { svc := newTestDeliveryService() taskID := "test-task-cancel-one" ch1, cancel1 := svc.SubscribeTask(taskID) _, cancel2 := svc.SubscribeTask(taskID) _ = cancel2 // 取消第一个订阅者 cancel1() // 验证还剩一个订阅者 svc.streamMu.Lock() subs := svc.streams[taskID] if len(subs) != 1 { svc.streamMu.Unlock() t.Fatalf("expected 1 subscriber after cancel, got %d", len(subs)) } svc.streamMu.Unlock() // 推送 snapshot snapshot := DeliveryTaskSnapshot{ Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning}, } svc.streamMu.Lock() for ch := range svc.streams[taskID] { select { case ch <- snapshot: default: } } svc.streamMu.Unlock() // ch1 已关闭,不应收到消息 select { case _, ok := <-ch1: if ok { t.Fatal("ch1 should be closed after cancel") } case <-time.After(100 * time.Millisecond): // OK: channel is closed, no value received } } func TestBroadcastTask_ClosedChannelCleanup(t *testing.T) { svc := newTestDeliveryService() taskID := "test-task-cleanup" // 创建一个订阅者并立即关闭 _, cancel := svc.SubscribeTask(taskID) cancel() // 等待 cancel 完成 time.Sleep(10 * time.Millisecond) // 创建一个新的正常订阅者 ch2, cancel2 := svc.SubscribeTask(taskID) defer cancel2() // 模拟 broadcastTask 行为(带 recover) snapshot := DeliveryTaskSnapshot{ Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning}, } svc.streamMu.Lock() var closedChannels []chan DeliveryTaskSnapshot for ch := range svc.streams[taskID] { func() { defer func() { if r := recover(); r != nil { closedChannels = append(closedChannels, ch) } }() select { case ch <- snapshot: default: } }() } for _, ch := range closedChannels { delete(svc.streams[taskID], ch) } if len(svc.streams[taskID]) == 0 { delete(svc.streams, taskID) } svc.streamMu.Unlock() // 验证 closed channel 被清理 svc.streamMu.Lock() if subs := svc.streams[taskID]; subs != nil && len(subs) != 1 { svc.streamMu.Unlock() t.Fatalf("expected 1 subscriber after cleanup, got %d", len(subs)) } svc.streamMu.Unlock() // 正常订阅者应该收到消息 select { case received := <-ch2: if received.Task == nil || received.Task.ID != taskID { t.Fatalf("expected task ID %q, got %+v", taskID, received) } case <-time.After(time.Second): t.Fatal("timeout waiting for snapshot from normal subscriber") } } func TestInvalidateStdoutCache(t *testing.T) { svc := newTestDeliveryService() // 填充缓存 svc.cacheMu.Lock() svc.stdoutCache["job-1"] = &stdoutCacheItem{stdout: "cached output", createdAt: time.Now()} svc.cacheMu.Unlock() // 验证缓存命中 svc.cacheMu.RLock() item, ok := svc.stdoutCache["job-1"] svc.cacheMu.RUnlock() if !ok { t.Fatal("cache entry not found before invalidation") } if item.stdout != "cached output" { t.Fatalf("cache stdout = %q, want %q", item.stdout, "cached output") } // 失效缓存 svc.InvalidateStdoutCache("job-1") // 验证缓存已失效 svc.cacheMu.RLock() _, ok = svc.stdoutCache["job-1"] svc.cacheMu.RUnlock() if ok { t.Fatal("cache entry still exists after invalidation") } } func TestInvalidateStdoutCache_NonExistent(t *testing.T) { svc := newTestDeliveryService() // 对不存在的 key 调用 invalidate 不应 panic svc.InvalidateStdoutCache("non-existent-job") // 验证缓存为空 svc.cacheMu.RLock() size := len(svc.stdoutCache) svc.cacheMu.RUnlock() if size != 0 { t.Fatalf("cache size = %d, want 0", size) } } func TestStdoutCache_TTLExpiry(t *testing.T) { svc := newTestDeliveryService() // 填充一个已过期的缓存条目 svc.cacheMu.Lock() svc.stdoutCache["job-expired"] = &stdoutCacheItem{ stdout: "old output", createdAt: time.Now().Add(-stdoutCacheTTL - time.Second), } svc.cacheMu.Unlock() // 模拟 AWXJobStdout 的缓存检查逻辑 svc.cacheMu.RLock() cacheHit := false if item, ok := svc.stdoutCache["job-expired"]; ok { if time.Since(item.createdAt) < stdoutCacheTTL { cacheHit = true } } svc.cacheMu.RUnlock() if cacheHit { t.Fatal("expired cache entry should not be a hit") } } func TestStdoutCache_FreshEntry(t *testing.T) { svc := newTestDeliveryService() // 填充一个新鲜的缓存条目 svc.cacheMu.Lock() svc.stdoutCache["job-fresh"] = &stdoutCacheItem{ stdout: "fresh output", createdAt: time.Now(), } svc.cacheMu.Unlock() // 模拟 AWXJobStdout 的缓存检查逻辑 svc.cacheMu.RLock() cacheHit := false var cachedStdout string if item, ok := svc.stdoutCache["job-fresh"]; ok { if time.Since(item.createdAt) < stdoutCacheTTL { cacheHit = true cachedStdout = item.stdout } } svc.cacheMu.RUnlock() if !cacheHit { t.Fatal("fresh cache entry should be a hit") } if cachedStdout != "fresh output" { t.Fatalf("cached stdout = %q, want %q", cachedStdout, "fresh output") } } func TestSubscribeTask_ConcurrentSafety(t *testing.T) { svc := newTestDeliveryService() taskID := "test-concurrent" var wg sync.WaitGroup const goroutines = 50 // 并发订阅 cancels := make([]func(), goroutines) for i := 0; i < goroutines; i++ { wg.Add(1) go func(idx int) { defer wg.Done() _, cancel := svc.SubscribeTask(taskID) cancels[idx] = cancel }(i) } wg.Wait() // 验证所有订阅者都已注册 svc.streamMu.Lock() subs := svc.streams[taskID] if len(subs) != goroutines { svc.streamMu.Unlock() t.Fatalf("expected %d subscribers, got %d", goroutines, len(subs)) } svc.streamMu.Unlock() // 并发取消 for _, cancel := range cancels { wg.Add(1) go func(c func()) { defer wg.Done() c() }(cancel) } wg.Wait() // 验证所有订阅者都已移除 svc.streamMu.Lock() if subs := svc.streams[taskID]; subs != nil && len(subs) > 0 { svc.streamMu.Unlock() t.Fatalf("expected 0 subscribers after concurrent cancel, got %d", len(subs)) } svc.streamMu.Unlock() } func TestStdoutCache_ConcurrentAccess(t *testing.T) { svc := newTestDeliveryService() var wg sync.WaitGroup const goroutines = 50 // 并发写入缓存 for i := 0; i < goroutines; i++ { wg.Add(1) go func(idx int) { defer wg.Done() jobID := "job-" + string(rune('A'+idx%26)) svc.cacheMu.Lock() svc.stdoutCache[jobID] = &stdoutCacheItem{ stdout: "output", createdAt: time.Now(), } svc.cacheMu.Unlock() }(i) } // 并发读取缓存 for i := 0; i < goroutines; i++ { wg.Add(1) go func(idx int) { defer wg.Done() jobID := "job-" + string(rune('A'+idx%26)) svc.cacheMu.RLock() _ = svc.stdoutCache[jobID] svc.cacheMu.RUnlock() }(i) } // 并发失效缓存 for i := 0; i < goroutines; i++ { wg.Add(1) go func(idx int) { defer wg.Done() jobID := "job-" + string(rune('A'+idx%26)) svc.InvalidateStdoutCache(jobID) }(i) } wg.Wait() }