WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/datasource/scheduler_test.go

@@ -0,0 +1,379 @@
+package datasource
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/hibiken/asynq"
+)
+
+// ──────────────────────────────────────────────────────────────────────
+// Fake implementations for testing
+// ──────────────────────────────────────────────────────────────────────
+
+// fakeDataSourceRepo is an in-memory DataSourceRepository.
+type fakeDataSourceRepo struct {
+	mu          sync.Mutex
+	dataSources map[string]*types.DataSource
+}
+
+func newFakeDataSourceRepo() *fakeDataSourceRepo {
+	return &fakeDataSourceRepo{dataSources: make(map[string]*types.DataSource)}
+}
+
+func (r *fakeDataSourceRepo) Create(_ context.Context, ds *types.DataSource) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.dataSources[ds.ID] = ds
+	return nil
+}
+
+func (r *fakeDataSourceRepo) FindByID(_ context.Context, id string) (*types.DataSource, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	ds, ok := r.dataSources[id]
+	if !ok {
+		return nil, ErrDataSourceNotFound
+	}
+	return ds, nil
+}
+
+func (r *fakeDataSourceRepo) FindByKnowledgeBase(_ context.Context, kbID string) ([]*types.DataSource, error) {
+	return nil, nil
+}
+
+func (r *fakeDataSourceRepo) Update(_ context.Context, ds *types.DataSource) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.dataSources[ds.ID] = ds
+	return nil
+}
+
+func (r *fakeDataSourceRepo) UpdateSyncState(ctx context.Context, ds *types.DataSource) error {
+	return r.Update(ctx, ds)
+}
+
+func (r *fakeDataSourceRepo) Delete(_ context.Context, id string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	delete(r.dataSources, id)
+	return nil
+}
+
+func (r *fakeDataSourceRepo) FindActive(_ context.Context) ([]*types.DataSource, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	var result []*types.DataSource
+	for _, ds := range r.dataSources {
+		if ds.Status == types.DataSourceStatusActive && ds.SyncSchedule != "" {
+			result = append(result, ds)
+		}
+	}
+	return result, nil
+}
+
+// fakeSyncLogRepo is an in-memory SyncLogRepository.
+type fakeSyncLogRepo struct {
+	mu   sync.Mutex
+	logs map[string]*types.SyncLog
+}
+
+func newFakeSyncLogRepo() *fakeSyncLogRepo {
+	return &fakeSyncLogRepo{logs: make(map[string]*types.SyncLog)}
+}
+
+func (r *fakeSyncLogRepo) Create(_ context.Context, log *types.SyncLog) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if log.ID == "" {
+		log.ID = "log-" + time.Now().Format("150405.000")
+	}
+	r.logs[log.ID] = log
+	return nil
+}
+
+func (r *fakeSyncLogRepo) FindByID(_ context.Context, id string) (*types.SyncLog, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	l, ok := r.logs[id]
+	if !ok {
+		return nil, ErrSyncLogNotFound
+	}
+	return l, nil
+}
+
+func (r *fakeSyncLogRepo) FindByDataSource(_ context.Context, dsID string, limit, offset int) ([]*types.SyncLog, error) {
+	return nil, nil
+}
+
+func (r *fakeSyncLogRepo) FindLatest(_ context.Context, dsID string) (*types.SyncLog, error) {
+	return nil, nil
+}
+
+func (r *fakeSyncLogRepo) Update(_ context.Context, log *types.SyncLog) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.logs[log.ID] = log
+	return nil
+}
+
+func (r *fakeSyncLogRepo) UpdateResult(ctx context.Context, log *types.SyncLog) error {
+	return r.Update(ctx, log)
+}
+
+func (r *fakeSyncLogRepo) CancelPendingByDataSource(_ context.Context, dsID string) error {
+	return nil
+}
+
+func (r *fakeSyncLogRepo) CleanupOldLogs(_ context.Context, retentionDays int) error {
+	return nil
+}
+
+func (r *fakeSyncLogRepo) HasRunningSync(_ context.Context, dsID string) (bool, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	for _, log := range r.logs {
+		if log.DataSourceID == dsID && log.Status == types.SyncLogStatusRunning {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+// fakeTaskEnqueuer counts how many tasks are enqueued.
+type fakeTaskEnqueuer struct {
+	count     atomic.Int64
+	lastQueue atomic.Value
+}
+
+func (e *fakeTaskEnqueuer) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
+	e.count.Add(1)
+	for _, opt := range opts {
+		if opt.Type() == asynq.QueueOpt {
+			if queue, ok := opt.Value().(string); ok {
+				e.lastQueue.Store(queue)
+			}
+		}
+	}
+	return &asynq.TaskInfo{ID: "task-fake"}, nil
+}
+
+// ──────────────────────────────────────────────────────────────────────
+// Tests
+// ──────────────────────────────────────────────────────────────────────
+
+func TestScheduler_StartWithActiveDataSources(t *testing.T) {
+	repo := newFakeDataSourceRepo()
+	_ = repo.Create(context.Background(), &types.DataSource{
+		ID:           "ds-1",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "*/2 * * * * *", // every 2 seconds (6-field cron with seconds)
+	})
+	_ = repo.Create(context.Background(), &types.DataSource{
+		ID:           "ds-2",
+		TenantID:     1,
+		Status:       types.DataSourceStatusPaused, // should NOT be scheduled
+		SyncSchedule: "*/2 * * * * *",
+	})
+
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
+
+	if err := scheduler.Start(context.Background()); err != nil {
+		t.Fatalf("Start() error: %v", err)
+	}
+	defer scheduler.Stop()
+
+	// Only ds-1 should be registered (ds-2 is paused, not returned by FindActive)
+	if scheduler.EntryCount() != 1 {
+		t.Errorf("EntryCount() = %d, want 1", scheduler.EntryCount())
+	}
+}
+
+func TestScheduler_CronFires(t *testing.T) {
+	repo := newFakeDataSourceRepo()
+	_ = repo.Create(context.Background(), &types.DataSource{
+		ID:           "ds-fire",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "* * * * * *", // every second
+	})
+
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
+
+	if err := scheduler.Start(context.Background()); err != nil {
+		t.Fatalf("Start() error: %v", err)
+	}
+
+	// Wait for at least one tick
+	time.Sleep(2500 * time.Millisecond)
+	scheduler.Stop()
+
+	if enqueuer.count.Load() == 0 {
+		t.Error("expected at least 1 enqueue, got 0")
+	}
+	if queue, _ := enqueuer.lastQueue.Load().(string); queue != types.QueueSync {
+		t.Errorf("scheduled sync queue = %q, want %q", queue, types.QueueSync)
+	}
+}
+
+func TestScheduler_AddOrUpdate(t *testing.T) {
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
+	scheduler.cron.Start()
+	defer scheduler.Stop()
+
+	ds := &types.DataSource{
+		ID:           "ds-new",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "0 0 * * * *", // every hour
+	}
+
+	// Add
+	if err := scheduler.AddOrUpdate(ds); err != nil {
+		t.Fatalf("AddOrUpdate() error: %v", err)
+	}
+	if scheduler.EntryCount() != 1 {
+		t.Errorf("EntryCount() = %d, want 1", scheduler.EntryCount())
+	}
+
+	// Update schedule
+	ds.SyncSchedule = "0 30 * * * *" // every half hour
+	if err := scheduler.AddOrUpdate(ds); err != nil {
+		t.Fatalf("AddOrUpdate() (update) error: %v", err)
+	}
+	if scheduler.EntryCount() != 1 {
+		t.Errorf("after update: EntryCount() = %d, want 1 (should replace, not add)", scheduler.EntryCount())
+	}
+}
+
+func TestScheduler_AddOrUpdate_PausedIsNoop(t *testing.T) {
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
+	scheduler.cron.Start()
+	defer scheduler.Stop()
+
+	ds := &types.DataSource{
+		ID:           "ds-paused",
+		TenantID:     1,
+		Status:       types.DataSourceStatusPaused,
+		SyncSchedule: "0 0 * * * *",
+	}
+
+	if err := scheduler.AddOrUpdate(ds); err != nil {
+		t.Fatalf("AddOrUpdate() error: %v", err)
+	}
+	if scheduler.EntryCount() != 0 {
+		t.Errorf("paused ds should not be scheduled, EntryCount() = %d", scheduler.EntryCount())
+	}
+}
+
+func TestScheduler_AddOrUpdate_EmptyScheduleIsNoop(t *testing.T) {
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
+	scheduler.cron.Start()
+	defer scheduler.Stop()
+
+	ds := &types.DataSource{
+		ID:           "ds-no-sched",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "",
+	}
+
+	if err := scheduler.AddOrUpdate(ds); err != nil {
+		t.Fatalf("AddOrUpdate() error: %v", err)
+	}
+	if scheduler.EntryCount() != 0 {
+		t.Errorf("empty schedule should not be registered, EntryCount() = %d", scheduler.EntryCount())
+	}
+}
+
+func TestScheduler_Remove(t *testing.T) {
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
+	scheduler.cron.Start()
+	defer scheduler.Stop()
+
+	ds := &types.DataSource{
+		ID:           "ds-rm",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "0 0 * * * *",
+	}
+
+	_ = scheduler.AddOrUpdate(ds)
+	if scheduler.EntryCount() != 1 {
+		t.Fatalf("pre-remove: EntryCount() = %d, want 1", scheduler.EntryCount())
+	}
+
+	scheduler.Remove("ds-rm")
+	if scheduler.EntryCount() != 0 {
+		t.Errorf("post-remove: EntryCount() = %d, want 0", scheduler.EntryCount())
+	}
+
+	// Remove non-existent is safe
+	scheduler.Remove("does-not-exist")
+}
+
+func TestScheduler_InvalidCron(t *testing.T) {
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(newFakeDataSourceRepo(), newFakeSyncLogRepo(), enqueuer)
+	scheduler.cron.Start()
+	defer scheduler.Stop()
+
+	ds := &types.DataSource{
+		ID:           "ds-bad",
+		TenantID:     1,
+		Status:       types.DataSourceStatusActive,
+		SyncSchedule: "not a cron",
+	}
+
+	err := scheduler.AddOrUpdate(ds)
+	if err == nil {
+		t.Fatal("expected error for invalid cron expression")
+	}
+	if scheduler.EntryCount() != 0 {
+		t.Errorf("invalid cron should not be registered, EntryCount() = %d", scheduler.EntryCount())
+	}
+}
+
+func TestScheduler_TriggerSync_InactiveSkipped(t *testing.T) {
+	repo := newFakeDataSourceRepo()
+	// Create a data source that is paused
+	_ = repo.Create(context.Background(), &types.DataSource{
+		ID:       "ds-inactive",
+		TenantID: 1,
+		Status:   types.DataSourceStatusPaused,
+	})
+
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
+
+	// Directly call triggerSync — it should skip because ds is not active
+	scheduler.triggerSync("ds-inactive", 1)
+
+	if enqueuer.count.Load() != 0 {
+		t.Error("should not enqueue for inactive data source")
+	}
+}
+
+func TestScheduler_TriggerSync_NotFound(t *testing.T) {
+	repo := newFakeDataSourceRepo()
+	enqueuer := &fakeTaskEnqueuer{}
+	scheduler := NewScheduler(repo, newFakeSyncLogRepo(), enqueuer)
+
+	// Should not panic, just skip
+	scheduler.triggerSync("nonexistent", 1)
+
+	if enqueuer.count.Load() != 0 {
+		t.Error("should not enqueue for non-existent data source")
+	}
+}