ragflow/internal/syncer/syncer_test.go

//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package syncer

import (
	"context"
	"errors"
	"io"
	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/entity"
	"ragflow/internal/service"
	syncerconnector "ragflow/internal/syncer/connector"
	connectormock "ragflow/internal/syncer/connector/mock"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

// fakeSink records document writes for coordinator tests.
type fakeSink struct {
	mu             sync.Mutex
	delay          time.Duration
	onUpsert       func(input service.DocumentUpsertInput)
	current        int
	maxConcurrent  int
	calls          []service.DocumentUpsertInput
	errBySourceID  map[string]error
	autoParseByDoc map[string]bool
}

// Upsert records one document write.
func (s *fakeSink) Upsert(ctx context.Context, input service.DocumentUpsertInput) (service.DocumentUpsertResult, error) {
	s.mu.Lock()
	s.current++
	if s.current > s.maxConcurrent {
		s.maxConcurrent = s.current
	}
	s.calls = append(s.calls, input)
	if s.autoParseByDoc == nil {
		s.autoParseByDoc = map[string]bool{}
	}
	s.autoParseByDoc[input.SourceDocument.SourceID] = input.AutoParse
	s.mu.Unlock()
	if s.onUpsert != nil {
		s.onUpsert(input)
	}
	if s.delay > 0 {
		select {
		case <-ctx.Done():
			return service.DocumentUpsertResult{}, ctx.Err()
		case <-time.After(s.delay):
		}
	}
	s.mu.Lock()
	s.current--
	s.mu.Unlock()
	if err := s.errBySourceID[input.SourceDocument.SourceID]; err != nil {
		return service.DocumentUpsertResult{}, err
	}
	return service.DocumentUpsertResult{DocID: input.DocumentID, Action: service.DocumentActionAdded}, nil
}

// callCount returns the number of sink calls.
func (s *fakeSink) callCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.calls)
}

// fakeStore provides document IDs and fingerprints.
type fakeStore struct {
	ids          map[string]struct{}
	fingerprints map[string]string
}

// ListIDs returns configured IDs.
func (s fakeStore) ListIDs(ctx context.Context, kbID, sourceType string) (map[string]struct{}, error) {
	if s.ids == nil {
		return map[string]struct{}{}, nil
	}
	return s.ids, nil
}

// GetFingerprintsByIDs returns configured fingerprints for requested IDs.
func (s fakeStore) GetFingerprintsByIDs(ctx context.Context, kbID, sourceType string, ids []string) (map[string]string, error) {
	result := make(map[string]string, len(ids))
	if s.fingerprints == nil {
		return result, nil
	}
	for _, id := range ids {
		if fingerprint, ok := s.fingerprints[id]; ok {
			result[id] = fingerprint
		}
	}
	return result, nil
}

// fakeDeleter records document deletes.
type fakeDeleter struct {
	mu      sync.Mutex
	deleted []string
}

type fakeTaskHandle struct {
	mu         sync.Mutex
	msg        common.TaskMessage
	acks       int
	nacks      int
	inProgress int
}

func (h *fakeTaskHandle) GetMessage() common.TaskMessage { return h.msg }
func (h *fakeTaskHandle) Ack() error {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.acks++
	return nil
}
func (h *fakeTaskHandle) Nack() error {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.nacks++
	return nil
}
func (h *fakeTaskHandle) InProgress() error {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.inProgress++
	return nil
}

func (h *fakeTaskHandle) counts() (int, int) {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.acks, h.nacks
}

type fakeSyncTaskBroker struct {
	published []string
	fetched   []common.TaskHandle
}

func (b *fakeSyncTaskBroker) InitSyncerStream() error   { return nil }
func (b *fakeSyncTaskBroker) InitSyncerConsumer() error { return nil }
func (b *fakeSyncTaskBroker) PublishSyncerTask(taskID string) error {
	b.published = append(b.published, taskID)
	return nil
}
func (b *fakeSyncTaskBroker) FetchSyncerTasks(batchSize int) ([]common.TaskHandle, error) {
	if len(b.fetched) == 0 {
		return nil, nil
	}
	if batchSize > len(b.fetched) {
		batchSize = len(b.fetched)
	}
	out := b.fetched[:batchSize]
	b.fetched = b.fetched[batchSize:]
	return out, nil
}

// DeleteDocument records one delete.
func (d *fakeDeleter) DeleteDocument(ctx context.Context, docID string) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.deleted = append(d.deleted, docID)
	return nil
}

// setupSyncerDB creates a SQLite test database.
func setupSyncerDB(t *testing.T) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err = db.AutoMigrate(&entity.Connector{}, &entity.Connector2Kb{}, &entity.Knowledgebase{}, &entity.SyncLogs{}, &entity.Document{}); err != nil {
		t.Fatalf("migrate sqlite: %v", err)
	}
	orig := dao.DB
	dao.DB = db
	t.Cleanup(func() { dao.DB = orig })
	return db
}

// insertTaskContext inserts one connector, mapping, KB, and sync task.
func insertTaskContext(t *testing.T, db *gorm.DB, connectorID, kbID, taskID, taskType string) {
	t.Helper()
	now := time.Now().Add(-time.Hour).Truncate(time.Second)
	ts := now.UnixMilli()
	fromBeginning := "0"
	var existingConnector int64
	if err := db.Model(&entity.Connector{}).Where("id = ?", connectorID).Count(&existingConnector).Error; err != nil {
		t.Fatalf("count connector: %v", err)
	}
	if existingConnector == 0 {
		if err := db.Create(&entity.Connector{
			ID:          connectorID,
			TenantID:    "tenant-1",
			Name:        connectorID,
			Source:      "mock",
			InputType:   "poll",
			Config:      entity.JSONMap{"sync_deleted_files": true},
			RefreshFreq: 0,
			PruneFreq:   0,
			TimeoutSecs: 60,
			Status:      dao.SyncStatusSchedule,
			BaseModel:   entity.BaseModel{UpdateDate: &now, UpdateTime: &ts},
		}).Error; err != nil {
			t.Fatalf("insert connector: %v", err)
		}
	}
	if err := db.Create(&entity.Knowledgebase{
		ID:           kbID,
		TenantID:     "tenant-1",
		Name:         kbID,
		EmbdID:       "embd",
		CreatedBy:    "tenant-1",
		ParserID:     "naive",
		ParserConfig: entity.JSONMap{},
	}).Error; err != nil {
		t.Fatalf("insert kb: %v", err)
	}
	if err := db.Create(&entity.Connector2Kb{ID: connectorID + kbID, ConnectorID: connectorID, KbID: kbID, AutoParse: "1"}).Error; err != nil {
		t.Fatalf("insert mapping: %v", err)
	}
	if err := db.Create(&entity.SyncLogs{
		ID:            taskID,
		ConnectorID:   connectorID,
		KbID:          kbID,
		TaskType:      taskType,
		Status:        dao.SyncStatusSchedule,
		FromBeginning: &fromBeginning,
		TimeStarted:   &now,
		ErrorMsg:      "",
		BaseModel:     entity.BaseModel{UpdateDate: &now, UpdateTime: &ts},
	}).Error; err != nil {
		t.Fatalf("insert task: %v", err)
	}
}

func insertKnowledgebaseMapping(t *testing.T, db *gorm.DB, connectorID, kbID string) {
	t.Helper()
	if err := db.Create(&entity.Knowledgebase{
		ID:           kbID,
		TenantID:     "tenant-1",
		Name:         kbID,
		EmbdID:       "embd",
		CreatedBy:    "tenant-1",
		ParserID:     "naive",
		ParserConfig: entity.JSONMap{},
	}).Error; err != nil {
		t.Fatalf("insert kb: %v", err)
	}
	if err := db.Create(&entity.Connector2Kb{ID: connectorID + kbID, ConnectorID: connectorID, KbID: kbID, AutoParse: "1"}).Error; err != nil {
		t.Fatalf("insert mapping: %v", err)
	}
}

func insertSyncLog(t *testing.T, db *gorm.DB, connectorID, kbID, taskID, taskType string) {
	t.Helper()
	now := time.Now().Add(-time.Hour).Truncate(time.Second)
	ts := now.UnixMilli()
	fromBeginning := "0"
	if err := db.Create(&entity.SyncLogs{
		ID:            taskID,
		ConnectorID:   connectorID,
		KbID:          kbID,
		TaskType:      taskType,
		Status:        dao.SyncStatusSchedule,
		FromBeginning: &fromBeginning,
		TimeStarted:   &now,
		ErrorMsg:      "",
		BaseModel:     entity.BaseModel{UpdateDate: &now, UpdateTime: &ts},
	}).Error; err != nil {
		t.Fatalf("insert task: %v", err)
	}
}

// newTestRegistry creates a mock connector registry.
func newTestRegistry(connectors map[string]*connectormock.Connector) *syncerconnector.Registry {
	registry := syncerconnector.NewRegistry()
	registry.Register("mock", func(ctx context.Context, taskContext any) (syncerconnector.Connector, error) {
		row := taskContext.(dao.SyncTaskContext)
		return connectors[row.Connector.ID], nil
	})
	return registry
}

// newCoordinator creates a test coordinator.
func newCoordinator(taskService *service.SyncTaskService, registry *syncerconnector.Registry, sink service.DocumentSink, pruneService *service.SyncPruneService, store service.DocumentStore) *TaskCoordinator {
	executor := NewSyncJobExecutor(SyncJobExecutorConfig{WorkerCount: 16})
	return NewTaskCoordinator(TaskCoordinatorConfig{ItemRetryCount: 1, ItemRetryBaseDelay: time.Millisecond}, taskService, registry, sink, pruneService, service.NewDocumentIDResolver(store), executor)
}

// TestClaimBlocksSameConnectorKBRunningTasks verifies DB-backed task mutual exclusion.
func TestClaimBlocksSameConnectorKBRunningTasks(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	insertSyncLog(t, db, "conn-1", "kb-1", "task-2", dao.TaskTypePrune)
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	claimed, err := taskService.Claim(t.Context(), "task-1")
	if err != nil {
		t.Fatalf("claim task-1: %v", err)
	}
	if !claimed {
		t.Fatalf("task-1 not claimed")
	}
	claimed, err = taskService.Claim(t.Context(), "task-2")
	if err != nil {
		t.Fatalf("claim task-2: %v", err)
	}
	if claimed {
		t.Fatalf("task-2 claimed while task-1 is running")
	}
}

// TestClaimAllowsSameConnectorDifferentKB verifies connector work can run for different KBs.
func TestClaimAllowsSameConnectorDifferentKB(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	insertKnowledgebaseMapping(t, db, "conn-1", "kb-2")
	insertSyncLog(t, db, "conn-1", "kb-2", "task-2", dao.TaskTypeSync)
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	for _, taskID := range []string{"task-1", "task-2"} {
		claimed, err := taskService.Claim(t.Context(), taskID)
		if err != nil {
			t.Fatalf("claim %s: %v", taskID, err)
		}
		if !claimed {
			t.Fatalf("%s not claimed", taskID)
		}
	}
}

// TestSchedulerRequiresBroker verifies JetStream is mandatory for the scheduler.
func TestSchedulerRequiresBroker(t *testing.T) {
	scheduler := NewScheduler(time.Hour, make(chan TaskEnvelope, 1), service.NewSyncTaskService(dao.NewSyncTaskDAO(nil)))
	err := scheduler.Run(t.Context())
	if err == nil || !strings.Contains(err.Error(), "NATS broker") {
		t.Fatalf("Run error = %v, want missing broker", err)
	}
}

// TestNATSSchedulerPublishesDueTasksWithoutClaiming verifies NATS mode leaves MySQL claim to workers.
func TestNATSSchedulerPublishesDueTasksWithoutClaiming(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	broker := &fakeSyncTaskBroker{}
	scheduler := NewNATSScheduler(time.Hour, make(chan TaskEnvelope, 1), taskService, broker)
	if err := scheduler.publishStartupTasks(t.Context()); err != nil {
		t.Fatalf("publish startup tasks: %v", err)
	}
	if len(broker.published) != 1 || broker.published[0] != "task-1" {
		t.Fatalf("published = %v", broker.published)
	}
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusSchedule {
		t.Fatalf("status = %s, want schedule", task.Status)
	}
}

// TestNATSSchedulerStartupPublishesScheduledTaskWithRefreshFreq verifies startup does not wait for the next refresh window.
func TestNATSSchedulerStartupPublishesScheduledTaskWithRefreshFreq(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	now := time.Now()
	updateTime := now.UnixMilli()
	if err := db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("refresh_freq", 5).Error; err != nil {
		t.Fatalf("set refresh freq: %v", err)
	}
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Updates(map[string]any{
		"update_date": now,
		"update_time": updateTime,
	}).Error; err != nil {
		t.Fatalf("set task update time: %v", err)
	}

	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	broker := &fakeSyncTaskBroker{}
	scheduler := NewNATSScheduler(time.Hour, make(chan TaskEnvelope, 1), taskService, broker)
	if err := scheduler.publishStartupTasks(t.Context()); err != nil {
		t.Fatalf("publish startup tasks: %v", err)
	}
	if len(broker.published) != 1 || broker.published[0] != "task-1" {
		t.Fatalf("published = %v", broker.published)
	}
}

// TestNATSSchedulerRecoversRunningTasksOnStartup verifies startup reconciliation publishes interrupted tasks.
func TestNATSSchedulerRecoversRunningTasksOnStartup(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error; err != nil {
		t.Fatalf("mark task running: %v", err)
	}
	if err := db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("status", dao.SyncStatusRunning).Error; err != nil {
		t.Fatalf("mark connector running: %v", err)
	}
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	broker := &fakeSyncTaskBroker{}
	scheduler := NewNATSScheduler(time.Hour, make(chan TaskEnvelope, 1), taskService, broker)

	if err := scheduler.publishStartupTasks(t.Context()); err != nil {
		t.Fatalf("publish startup tasks: %v", err)
	}
	if len(broker.published) != 1 || broker.published[0] != "task-1" {
		t.Fatalf("published = %v", broker.published)
	}
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusSchedule {
		t.Fatalf("task status = %s, want schedule", task.Status)
	}
	var connector entity.Connector
	if err := db.First(&connector, "id = ?", "conn-1").Error; err != nil {
		t.Fatalf("load connector: %v", err)
	}
	if connector.Status != dao.SyncStatusSchedule {
		t.Fatalf("connector status = %s, want schedule", connector.Status)
	}
}

// TestNATSSchedulerPublishesOnlyDueTasks verifies periodic NATS publishing respects refresh windows.
func TestNATSSchedulerPublishesOnlyDueTasks(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	now := time.Now()
	updateTime := now.UnixMilli()
	if err := db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("refresh_freq", 5).Error; err != nil {
		t.Fatalf("set refresh freq: %v", err)
	}
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Updates(map[string]any{
		"update_date": now,
		"update_time": updateTime,
	}).Error; err != nil {
		t.Fatalf("set fresh task update time: %v", err)
	}
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	broker := &fakeSyncTaskBroker{}
	scheduler := NewNATSScheduler(10*time.Millisecond, make(chan TaskEnvelope, 1), taskService, broker)

	if err := scheduler.publishDueTasks(t.Context()); err != nil {
		t.Fatalf("publish fresh due tasks: %v", err)
	}
	if len(broker.published) != 0 {
		t.Fatalf("fresh task published = %v, want none", broker.published)
	}

	dueAt := now.Add(-10 * time.Minute)
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Updates(map[string]any{
		"update_date": dueAt,
		"update_time": dueAt.UnixMilli(),
	}).Error; err != nil {
		t.Fatalf("set due task update time: %v", err)
	}

	if err := scheduler.publishDueTasks(t.Context()); err != nil {
		t.Fatalf("publish due tasks: %v", err)
	}
	if len(broker.published) != 1 || broker.published[0] != "task-1" {
		t.Fatalf("published = %v, want due task", broker.published)
	}
}

// TestNATSSchedulerBuffersFetchedTasks verifies NATS mode stores excess tasks in the local queue.
func TestNATSSchedulerBuffersFetchedTasks(t *testing.T) {
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(nil))
	queue := make(chan TaskEnvelope, 2)
	scheduler := NewNATSScheduler(time.Hour, queue, taskService, &fakeSyncTaskBroker{})
	handles := []common.TaskHandle{
		&fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-1", TaskType: common.TaskTypeSyncer}},
		&fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-2", TaskType: common.TaskTypeSyncer}},
	}

	if err := scheduler.enqueueHandles(t.Context(), handles); err != nil {
		t.Fatalf("enqueue handles: %v", err)
	}
	if got := scheduler.queueAvailable(); got != 0 {
		t.Fatalf("queue capacity while buffered = %d, want 0", got)
	}

	first := <-queue
	stopEnvelopeHeartbeat(first)
	if got := scheduler.queueAvailable(); got != 1 {
		t.Fatalf("queue capacity after dequeue = %d, want 1", got)
	}
	second := <-queue
	stopEnvelopeHeartbeat(second)
}

// TestWorkersRunDifferentConnectorsInParallel verifies task-level parallelism.
func TestWorkersRunDifferentConnectorsInParallel(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	insertTaskContext(t, db, "conn-2", "kb-2", "task-2", dao.TaskTypeSync)
	taskDAO := dao.NewSyncTaskDAO(db)
	taskService := service.NewSyncTaskService(taskDAO)
	now := time.Now()
	for _, id := range []string{"task-1", "task-2"} {
		claimed, err := taskDAO.ClaimTask(t.Context(), id, now)
		if err != nil || !claimed {
			t.Fatalf("claim %s: %v %v", id, claimed, err)
		}
	}
	sink := &fakeSink{delay: 120 * time.Millisecond}
	connectors := map[string]*connectormock.Connector{
		"conn-1": {SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: now}}}}},
		"conn-2": {SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "b", UpdatedAt: now}}}}},
	}
	queue := make(chan TaskEnvelope, 2)
	worker := NewTaskWorker(queue, taskService, newCoordinator(taskService, newTestRegistry(connectors), sink, nil, fakeStore{}), NewConnectorLock())
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	go worker.Run(ctx, 2)
	queue <- TaskEnvelope{TaskID: "task-1"}
	queue <- TaskEnvelope{TaskID: "task-2"}
	time.Sleep(40 * time.Millisecond)
	sink.mu.Lock()
	maxConcurrent := sink.maxConcurrent
	sink.mu.Unlock()
	if maxConcurrent < 2 {
		t.Fatalf("max concurrent sink calls = %d, want >= 2", maxConcurrent)
	}
}

// TestNATSTaskWorkerClaimsAndAcksOnSuccess verifies NATS messages are ACKed after durable completion.
func TestNATSTaskWorkerClaimsAndAcksOnSuccess(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	now := time.Now()
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: now}}}}}
	handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-1", TaskType: common.TaskTypeSyncer}}
	worker := NewTaskWorker(
		make(chan TaskEnvelope, 1),
		taskService,
		newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), &fakeSink{}, nil, fakeStore{}),
		NewConnectorLock(),
	)
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1", Handle: handle})

	if handle.acks != 1 || handle.nacks != 0 {
		t.Fatalf("settlement acks=%d nacks=%d", handle.acks, handle.nacks)
	}
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusDone {
		t.Fatalf("status = %s, want done", task.Status)
	}
}

// TestNATSTaskWorkerAcksUnclaimableMessage verifies duplicate/stale messages do not redeliver forever.
func TestNATSTaskWorkerAcksUnclaimableMessage(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error; err != nil {
		t.Fatalf("mark running: %v", err)
	}
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-1", TaskType: common.TaskTypeSyncer}}
	worker := NewTaskWorker(make(chan TaskEnvelope, 1), taskService, newCoordinator(taskService, newTestRegistry(nil), &fakeSink{}, nil, fakeStore{}), NewConnectorLock())
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1", Handle: handle})
	if handle.acks != 1 || handle.nacks != 0 {
		t.Fatalf("settlement acks=%d nacks=%d", handle.acks, handle.nacks)
	}
}

// TestSameConnectorDifferentKBsRunInParallel verifies one datasource can sync into different KBs concurrently.
func TestSameConnectorDifferentKBsRunInParallel(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	insertTaskContext(t, db, "conn-1", "kb-2", "task-2", dao.TaskTypeSync)
	taskDAO := dao.NewSyncTaskDAO(db)
	taskService := service.NewSyncTaskService(taskDAO)
	for _, id := range []string{"task-1", "task-2"} {
		if _, err := taskDAO.ClaimTask(t.Context(), id, time.Now()); err != nil {
			t.Fatalf("claim %s: %v", id, err)
		}
	}
	sink := &fakeSink{delay: 100 * time.Millisecond}
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: time.Now()}}}}}
	queue := make(chan TaskEnvelope, 2)
	worker := NewTaskWorker(queue, taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{}), NewConnectorLock())
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	go worker.Run(ctx, 2)
	queue <- TaskEnvelope{TaskID: "task-1"}
	queue <- TaskEnvelope{TaskID: "task-2"}
	time.Sleep(40 * time.Millisecond)
	sink.mu.Lock()
	maxConcurrent := sink.maxConcurrent
	sink.mu.Unlock()
	if maxConcurrent < 2 {
		t.Fatalf("max concurrent same-connector different-kb sink calls = %d, want >= 2", maxConcurrent)
	}
}

// TestConnectorKBLockSerializesSyncAndPrune verifies sync and prune for the same connector/KB do not overlap.
func TestConnectorKBLockSerializesSyncAndPrune(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	now := time.Now().Add(-time.Hour).Truncate(time.Second)
	ts := now.UnixMilli()
	fromBeginning := "0"
	if err := db.Create(&entity.SyncLogs{
		ID:            "task-2",
		ConnectorID:   "conn-1",
		KbID:          "kb-1",
		TaskType:      dao.TaskTypePrune,
		Status:        dao.SyncStatusSchedule,
		FromBeginning: &fromBeginning,
		TimeStarted:   &now,
		ErrorMsg:      "",
		BaseModel:     entity.BaseModel{UpdateDate: &now, UpdateTime: &ts},
	}).Error; err != nil {
		t.Fatalf("insert prune task: %v", err)
	}
	taskDAO := dao.NewSyncTaskDAO(db)
	taskService := service.NewSyncTaskService(taskDAO)
	for _, id := range []string{"task-1", "task-2"} {
		if _, err := taskDAO.ClaimTask(t.Context(), id, time.Now()); err != nil {
			t.Fatalf("claim %s: %v", id, err)
		}
	}
	sink := &fakeSink{delay: 120 * time.Millisecond}
	connector := &connectormock.Connector{
		SyncBatches:  []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: time.Now()}}}},
		PruneBatches: []syncerconnector.PruneBatch{{Documents: []syncerconnector.SlimDocument{{SourceID: "a"}}}},
	}
	pruneService := service.NewSyncPruneService(&fakeDeleter{}, fakeStore{ids: map[string]struct{}{}})
	worker := NewTaskWorker(make(chan TaskEnvelope, 2), taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, pruneService, fakeStore{}), NewConnectorLock())
	done := make(chan struct{})
	go func() {
		defer close(done)
		worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1"})
	}()
	waitForSinkConcurrency(t, sink, 1)
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-2"})
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-2").Error; err != nil {
		t.Fatalf("load prune task: %v", err)
	}
	if task.Status != dao.SyncStatusSchedule {
		t.Fatalf("same connector/kb prune status = %s, want schedule", task.Status)
	}
	<-done
}

func waitForSinkConcurrency(t *testing.T, sink *fakeSink, want int) {
	t.Helper()
	deadline := time.After(time.Second)
	ticker := time.NewTicker(time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-deadline:
			t.Fatalf("timed out waiting for sink concurrency >= %d", want)
		case <-ticker.C:
			sink.mu.Lock()
			current := sink.current
			sink.mu.Unlock()
			if current >= want {
				return
			}
		}
	}
}

// TestSyncRunnerSubmitsBatchesBeforeWaiting verifies source reads are not blocked by prior batch jobs.
func TestSyncRunnerSubmitsBatchesBeforeWaiting(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	taskDAO := dao.NewSyncTaskDAO(db)
	taskService := service.NewSyncTaskService(taskDAO)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	start := time.Now()
	var secondBatchAt time.Time
	connector := &connectormock.Connector{
		SyncBatches: []syncerconnector.SyncBatch{
			{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: start}, {SourceID: "b", UpdatedAt: start}}},
			{Documents: []syncerconnector.SourceDocument{{SourceID: "c", UpdatedAt: start}}},
		},
		OnSyncBatch: func(index int) {
			if index == 1 {
				secondBatchAt = time.Now()
			}
		},
	}
	sink := &fakeSink{delay: 80 * time.Millisecond}
	coordinator := newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{})
	taskContext, _ := taskService.GetContext(t.Context(), "task-1")
	if err := coordinator.Execute(t.Context(), taskContext, testLockLease()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	if secondBatchAt.Sub(start) >= 70*time.Millisecond {
		t.Fatalf("second batch was blocked by first batch for %s", secondBatchAt.Sub(start))
	}
}

// TestSyncRunnerProcessesBatchJobsInParallel verifies source batches run as parallel BatchJobs.
func TestSyncRunnerProcessesBatchJobsInParallel(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	now := time.Now()
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{
		{Documents: []syncerconnector.SourceDocument{{SourceID: "a", UpdatedAt: now}}},
		{Documents: []syncerconnector.SourceDocument{{SourceID: "b", UpdatedAt: now}}},
		{Documents: []syncerconnector.SourceDocument{{SourceID: "c", UpdatedAt: now}}},
	}}
	sink := &fakeSink{delay: 80 * time.Millisecond}
	coordinator := newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{})
	taskContext, _ := taskService.GetContext(t.Context(), "task-1")
	if err := coordinator.Execute(t.Context(), taskContext, testLockLease()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	sink.mu.Lock()
	maxConcurrent := sink.maxConcurrent
	sink.mu.Unlock()
	if maxConcurrent < 2 {
		t.Fatalf("max concurrent batches = %d, want >= 2", maxConcurrent)
	}
}

// TestHash128MatchesPythonGolden verifies Python-compatible xxhash128.
func TestHash128MatchesPythonGolden(t *testing.T) {
	cases := map[string]string{
		"connector-1:source-1":      "0b66c92e6cefff918067e3c42606cda8",
		"kb-1:connector-1:source-1": "ef785a8871b90d910f42cf31a8155476",
		"":                          "99aa06d3014798d86001c324468d497f",
		"hello":                     "b5e9c1ad071b3e7fc779cfaa5e523818",
	}
	for input, want := range cases {
		if got := service.Hash128(input); got != want {
			t.Fatalf("Hash128(%q) = %s, want %s", input, got, want)
		}
	}
}

// TestFingerprintSkipsUnchangedDocument verifies unchanged docs are not fetched or upserted.
func TestFingerprintSkipsUnchangedDocument(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	legacyID := service.Hash128("conn-1:source-1")
	store := fakeStore{ids: map[string]struct{}{legacyID: {}}, fingerprints: map[string]string{legacyID: "fp-1"}}
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "source-1", Fingerprint: "fp-1", FetchRef: &syncerconnector.FetchReference{Key: "lazy"}, UpdatedAt: time.Now()}}}}}
	sink := &fakeSink{}
	coordinator := newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, store)
	taskContext, _ := taskService.GetContext(t.Context(), "task-1")
	if err := coordinator.Execute(t.Context(), taskContext, testLockLease()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	if sink.callCount() != 0 {
		t.Fatalf("sink calls = %d, want 0", sink.callCount())
	}
}

// TestAutoParseFlagFlowsToSink verifies connector2kb.auto_parse is preserved.
func TestAutoParseFlagFlowsToSink(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	if err := db.Model(&entity.Connector2Kb{}).Where("connector_id = ? AND kb_id = ?", "conn-1", "kb-1").Update("auto_parse", "0").Error; err != nil {
		t.Fatalf("disable auto_parse: %v", err)
	}
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "source-1", Blob: []byte("x"), UpdatedAt: time.Now()}}}}}
	sink := &fakeSink{}
	coordinator := newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{})
	taskContext, _ := taskService.GetContext(t.Context(), "task-1")
	if err := coordinator.Execute(t.Context(), taskContext, testLockLease()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	if sink.autoParseByDoc["source-1"] {
		t.Fatalf("auto_parse flowed as true, want false")
	}
}

// TestCompleteSyncSchedulesNextRun verifies completing one sync keeps the connector schedulable.
func TestCompleteSyncSchedulesNextRun(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	_ = db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	now := time.Now()
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "source-1", Blob: []byte("x"), UpdatedAt: now}}}}}
	coordinator := newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), &fakeSink{}, nil, fakeStore{})
	taskContext, _ := taskService.GetContext(t.Context(), "task-1")
	if err := coordinator.Execute(t.Context(), taskContext, testLockLease()); err != nil {
		t.Fatalf("execute: %v", err)
	}

	var connectorRow entity.Connector
	if err := db.First(&connectorRow, "id = ?", "conn-1").Error; err != nil {
		t.Fatalf("load connector: %v", err)
	}
	if connectorRow.Status != dao.SyncStatusSchedule {
		t.Fatalf("connector status = %s, want schedule", connectorRow.Status)
	}
	var scheduled int64
	if err := db.Model(&entity.SyncLogs{}).
		Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", "conn-1", "kb-1", dao.TaskTypeSync, dao.SyncStatusSchedule).
		Count(&scheduled).Error; err != nil {
		t.Fatalf("count scheduled sync logs: %v", err)
	}
	if scheduled != 1 {
		t.Fatalf("scheduled sync logs = %d, want 1", scheduled)
	}
}

// TestCancelStopsRunningSync verifies a stop request prevents further work and completion.
func TestCancelStopsRunningSync(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	_ = db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	now := time.Now()
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{
		{SourceID: "source-1", Blob: []byte("one"), UpdatedAt: now},
		{SourceID: "source-2", Blob: []byte("two"), UpdatedAt: now},
	}}}}
	var cancelOnce sync.Once
	sink := &fakeSink{
		onUpsert: func(input service.DocumentUpsertInput) {
			cancelOnce.Do(func() {
				if err := db.Model(&entity.SyncLogs{}).Where("id = ?", input.TaskContext.Task.ID).Update("status", dao.SyncStatusCancel).Error; err != nil {
					t.Errorf("cancel sync log: %v", err)
				}
				if err := db.Model(&entity.Connector{}).Where("id = ?", input.TaskContext.Connector.ID).Update("status", dao.SyncStatusCancel).Error; err != nil {
					t.Errorf("cancel connector: %v", err)
				}
				time.Sleep(syncCancelCheckInterval)
			})
		},
	}
	worker := NewTaskWorker(make(chan TaskEnvelope, 1), taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{}), NewConnectorLock())
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1"})

	if calls := sink.callCount(); calls != 1 {
		t.Fatalf("sink calls = %d, want 1", calls)
	}
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusCancel {
		t.Fatalf("task status = %s, want cancel", task.Status)
	}
	var scheduled int64
	if err := db.Model(&entity.SyncLogs{}).
		Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", "conn-1", "kb-1", dao.TaskTypeSync, dao.SyncStatusSchedule).
		Count(&scheduled).Error; err != nil {
		t.Fatalf("count scheduled sync logs: %v", err)
	}
	if scheduled != 0 {
		t.Fatalf("scheduled sync logs = %d, want 0", scheduled)
	}
	var connectorRow entity.Connector
	if err := db.First(&connectorRow, "id = ?", "conn-1").Error; err != nil {
		t.Fatalf("load connector: %v", err)
	}
	if connectorRow.Status != dao.SyncStatusCancel {
		t.Fatalf("connector status = %s, want cancel", connectorRow.Status)
	}
}

// TestSyncRunnerResultWaitHonorsCancel verifies Run does not block forever waiting for job results after cancellation.
func TestSyncRunnerResultWaitHonorsCancel(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	taskContext, err := taskService.GetContext(t.Context(), "task-1")
	if err != nil {
		t.Fatalf("get context: %v", err)
	}
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "source-1", Blob: []byte("x"), UpdatedAt: time.Now()}}}}}
	queue := &SyncJobQueue{taskID: "task-1", jobs: make(chan *syncJob, 1)}
	runner := NewSyncRunner(TaskCoordinatorConfig{ItemRetryCount: 1, ItemRetryBaseDelay: time.Millisecond}, taskService, &fakeSink{}, service.NewDocumentIDResolver(fakeStore{}), queue)

	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() {
		result <- runner.Run(ctx, taskContext, connector)
	}()

	select {
	case <-queue.jobs:
	case <-time.After(time.Second):
		t.Fatalf("timed out waiting for submitted job")
	}
	cancel()

	select {
	case err = <-result:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("Run error = %v, want context canceled", err)
		}
	case <-time.After(time.Second):
		t.Fatalf("Run did not return after cancellation")
	}
}

// TestBatchFailureDoesNotAdvanceWaterline verifies failed tasks keep poll_range_end.
func TestBatchFailureDoesNotAdvanceWaterline(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	connector := &connectormock.Connector{SyncBatches: []syncerconnector.SyncBatch{{Documents: []syncerconnector.SourceDocument{{SourceID: "bad", Blob: []byte("x"), UpdatedAt: time.Now()}}}}}
	sink := &fakeSink{errBySourceID: map[string]error{"bad": errors.New("boom")}}
	queue := make(chan TaskEnvelope, 1)
	worker := NewTaskWorker(queue, taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), sink, nil, fakeStore{}), NewConnectorLock())
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1"})
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusFail {
		t.Fatalf("status = %s, want fail", task.Status)
	}
	if task.PollRangeEnd != nil {
		t.Fatalf("poll_range_end advanced on failure")
	}
}

// TestRecoverStaleRunningTasks verifies timeout recovery.
func TestRecoverStaleRunningTasks(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
	started := time.Now().Add(-2 * time.Hour)
	if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Updates(map[string]any{"status": dao.SyncStatusRunning, "time_started": started}).Error; err != nil {
		t.Fatalf("mark running: %v", err)
	}
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	if err := taskService.RecoverStaleRunning(t.Context(), time.Now()); err != nil {
		t.Fatalf("recover: %v", err)
	}
	var task entity.SyncLogs
	if err := db.First(&task, "id = ?", "task-1").Error; err != nil {
		t.Fatalf("load task: %v", err)
	}
	if task.Status != dao.SyncStatusSchedule {
		t.Fatalf("status = %s, want schedule", task.Status)
	}
}

// TestPruneSourceFailureDoesNotDelete verifies incomplete source listings never delete.
func TestPruneSourceFailureDoesNotDelete(t *testing.T) {
	db := setupSyncerDB(t)
	insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypePrune)
	_ = db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error
	taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
	deleter := &fakeDeleter{}
	pruneService := service.NewSyncPruneService(deleter, fakeStore{ids: map[string]struct{}{"stale": {}}})
	connector := &connectormock.Connector{PruneErrAt: 1, PruneBatches: []syncerconnector.PruneBatch{{Documents: []syncerconnector.SlimDocument{{SourceID: "keep"}}}}}
	queue := make(chan TaskEnvelope, 1)
	worker := NewTaskWorker(queue, taskService, newCoordinator(taskService, newTestRegistry(map[string]*connectormock.Connector{"conn-1": connector}), &fakeSink{}, pruneService, fakeStore{}), NewConnectorLock())
	worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1"})
	if len(deleter.deleted) != 0 {
		t.Fatalf("deleted %v despite incomplete snapshot", deleter.deleted)
	}
}

// TestPruneRetainIDsIncludeLegacyAndNew verifies PRUNE retains both ID schemes.
func TestPruneRetainIDsIncludeLegacyAndNew(t *testing.T) {
	retain := service.RetainDocumentIDs("kb-1", "conn-1", []string{"source-1"})
	legacy := service.Hash128("conn-1:source-1")
	next := service.Hash128("kb-1:conn-1:source-1")
	if _, ok := retain[legacy]; !ok {
		t.Fatalf("legacy ID missing from retain set")
	}
	if _, ok := retain[next]; !ok {
		t.Fatalf("new ID missing from retain set")
	}
}

// TestMockSessionEOF documents the mock session EOF contract.
func TestMockSessionEOF(t *testing.T) {
	session, _ := (&connectormock.Connector{}).OpenSync(t.Context(), syncerconnector.SyncRequest{})
	_, err := session.NextBatch(t.Context())
	if !errors.Is(err, io.EOF) {
		t.Fatalf("err = %v, want io.EOF", err)
	}
}

// TestTaskExecutionDeadlineCapsConnectorTimeout verifies tasks cannot outlive the connector lock lease.
func TestTaskExecutionDeadlineCapsConnectorTimeout(t *testing.T) {
	now := time.Date(2026, 8, 10, 10, 0, 0, 0, time.UTC)
	taskContext := service.SyncTaskContext{}
	taskContext.Connector.TimeoutSecs = int64(connectorLockTTL.Seconds()) + 60
	lease := ConnectorLockLease{ExpiresAt: now.Add(10 * time.Minute)}
	if got, want := taskExecutionDeadline(now, taskContext, lease), lease.ExpiresAt.Add(-connectorLockSafetyMargin); !got.Equal(want) {
		t.Fatalf("deadline = %s, want %s", got, want)
	}
	taskContext.Connector.TimeoutSecs = 30
	if got, want := taskExecutionDeadline(now, taskContext, lease), now.Add(30*time.Second); !got.Equal(want) {
		t.Fatalf("deadline = %s, want %s", got, want)
	}
	lease.ExpiresAt = now.Add(time.Second)
	if got := taskExecutionDeadline(now, taskContext, lease); !got.Equal(now) {
		t.Fatalf("deadline = %s, want %s", got, now)
	}
}

func testLockLease() ConnectorLockLease {
	return ConnectorLockLease{ExpiresAt: time.Now().Add(connectorLockTTL)}
}