alt/services/worker/internal/scheduler/runner_test.go

package scheduler

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"git.toki-labs.com/toki/alt/services/worker/internal/jobs"
	"git.toki-labs.com/toki/alt/services/worker/internal/marketdata/importer"
)

type fakeExecutor struct {
	mu           sync.Mutex
	activeCount  int
	maxActive    int
	executedJobs []jobs.Job
	executeFn    func(ctx context.Context, job jobs.Job) error
	// resultFn controls ExecuteWithResult returns.
	// nil = return (Result{}, false, nil) — legacy no-result path.
	resultFn func(ctx context.Context, job jobs.Job) (jobs.Result, bool, error)
}

func (f *fakeExecutor) Execute(ctx context.Context, job jobs.Job) error {
	f.mu.Lock()
	f.activeCount++
	if f.activeCount > f.maxActive {
		f.maxActive = f.activeCount
	}
	f.executedJobs = append(f.executedJobs, job)
	f.mu.Unlock()

	defer func() {
		f.mu.Lock()
		f.activeCount--
		f.mu.Unlock()
	}()

	if f.executeFn != nil {
		return f.executeFn(ctx, job)
	}
	return nil
}

func (f *fakeExecutor) ExecuteWithResult(ctx context.Context, job jobs.Job) (jobs.Result, bool, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.resultFn != nil {
		return f.resultFn(ctx, job)
	}
	// Default: legacy no-result path.
	return jobs.Result{}, false, nil
}

func TestRunTickDispatchesItemsWithBoundedParallelism(t *testing.T) {
	resetActiveJobs()
	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "kr-daily-parallel",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"A", "B", "C", "D"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 2,
			},
		},
	}

	exec := &fakeExecutor{
		executeFn: func(ctx context.Context, job jobs.Job) error {
			time.Sleep(10 * time.Millisecond)
			return nil
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec,
	}

	res, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("RunTick: %v", err)
	}

	if len(res.Items) != 4 {
		t.Errorf("expected 4 item results, got %d", len(res.Items))
	}

	if exec.maxActive > 2 {
		t.Errorf("max concurrent jobs %d exceeded parallelism limit 2", exec.maxActive)
	}

	for _, item := range res.Items {
		if item.Status != "succeeded" {
			t.Errorf("expected succeeded status, got %s", item.Status)
		}
	}
}

func TestRunTickSkipsDuplicateRunningWindow(t *testing.T) {
	resetActiveJobs()
	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "kr-daily-dup",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"A"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 1,
			},
		},
	}

	startChan := make(chan struct{})
	doneChan := make(chan struct{})

	exec := &fakeExecutor{
		executeFn: func(ctx context.Context, job jobs.Job) error {
			close(startChan)
			<-doneChan
			return nil
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec,
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		_, _ = RunTick(context.Background(), cfg, opts)
	}()

	<-startChan // Wait until first tick starts executing

	res2, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("second RunTick: %v", err)
	}

	if len(res2.Items) != 1 {
		t.Fatalf("expected 1 result in duplicate run, got %d", len(res2.Items))
	}
	if res2.Items[0].Status != "skipped" {
		t.Errorf("expected duplicate run to be skipped, got %s", res2.Items[0].Status)
	}
	if res2.Items[0].Error != "duplicate window running" {
		t.Errorf("expected error string duplicate window running, got %q", res2.Items[0].Error)
	}

	close(doneChan)
	wg.Wait()

	// Sequential 2nd run
	exec3 := &fakeExecutor{}
	opts3 := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec3,
	}
	res3, err := RunTick(context.Background(), cfg, opts3)
	if err != nil {
		t.Fatalf("third sequential RunTick: %v", err)
	}
	if len(res3.Items) != 1 {
		t.Fatalf("expected 1 result, got %d", len(res3.Items))
	}
	if res3.Items[0].Status != "succeeded" {
		t.Errorf("expected succeeded status, got %s", res3.Items[0].Status)
	}
}

func TestWriteTickJSONLIncludesReadinessFields(t *testing.T) {
	res := TickResult{
		Items: []ItemResult{
			{
				Schedule:  "test-sched",
				Symbol:    "AAPL",
				Status:    "succeeded",
				Provider:  "kis",
				Timeframe: "daily",
				Readiness: "ready",
				Freshness: "fresh",
			},
			{
				Schedule:  "test-sched",
				Symbol:    "MSFT",
				Status:    "failed",
				Provider:  "kis",
				Timeframe: "daily",
				Readiness: "error",
				Freshness: "stale",
				Error:     "network error",
			},
		},
	}

	var jsonl bytes.Buffer
	if err := WriteTickJSONL(&jsonl, res); err != nil {
		t.Fatalf("WriteTickJSONL: %v", err)
	}

	lines := strings.Split(strings.TrimSpace(jsonl.String()), "\n")
	if len(lines) != 2 {
		t.Fatalf("expected 2 JSONL lines, got %d", len(lines))
	}

	if !strings.Contains(lines[0], `"readiness":"ready"`) || !strings.Contains(lines[0], `"freshness":"fresh"`) {
		t.Errorf("line 0 missing readiness/freshness fields: %s", lines[0])
	}
	if !strings.Contains(lines[1], `"readiness":"error"`) || !strings.Contains(lines[1], `"freshness":"stale"`) {
		t.Errorf("line 1 missing readiness/freshness fields: %s", lines[1])
	}

	var text bytes.Buffer
	if err := WriteTickText(&text, res); err != nil {
		t.Fatalf("WriteTickText: %v", err)
	}
	textLines := strings.Split(strings.TrimSpace(text.String()), "\n")
	if len(textLines) != 2 {
		t.Fatalf("expected 2 text lines, got %d", len(textLines))
	}
	if !strings.Contains(textLines[0], "readiness=ready freshness=fresh") {
		t.Errorf("text line 0 missing readiness/freshness: %s", textLines[0])
	}
	if !strings.Contains(textLines[1], "readiness=error freshness=stale") {
		t.Errorf("text line 1 missing readiness/freshness: %s", textLines[1])
	}
}

type fakeDailyBarImporter struct {
	mu           sync.Mutex
	upsertedBars map[string]int
}

func (f *fakeDailyBarImporter) ImportDailyBars(ctx context.Context, request importer.DailyBarRequest) (importer.Result, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	timeframe := "daily"
	days := int(request.To.Sub(request.From).Hours()/24) + 1
	for _, sym := range request.Selector.Symbols {
		for i := 0; i < days; i++ {
			t := request.From.Add(time.Duration(i) * 24 * time.Hour)
			dateStr := t.Format("2006-01-02")
			key := fmt.Sprintf("%s:%s:%s", sym, timeframe, dateStr)
			f.upsertedBars[key]++
		}
	}
	return importer.Result{
		Instruments: len(request.Selector.Symbols),
		Bars:        len(request.Selector.Symbols) * days,
	}, nil
}

func TestRunTickSequentialTicksDoNotDuplicateBarKeys(t *testing.T) {
	resetActiveJobs()

	fakeImp := &fakeDailyBarImporter{
		upsertedBars: make(map[string]int),
	}

	runner := jobs.NewRunner()
	jobs.RegisterBuiltins(runner)
	cap := testCapability()
	jobs.RegisterDailyBarImportHandler(runner, cap, fakeImp)

	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "kr-daily-dup-assert",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"005930"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 1,
			},
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(cap),
		Executor:     runner,
	}

	// 1st tick
	res1, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("first tick: %v", err)
	}
	if len(res1.Items) != 1 || res1.Items[0].Status != "succeeded" {
		t.Fatalf("expected succeeded in 1st tick, got: %+v", res1.Items)
	}

	fakeImp.mu.Lock()
	keysCount1 := len(fakeImp.upsertedBars)
	if keysCount1 == 0 {
		fakeImp.mu.Unlock()
		t.Fatal("expected upserted bars count > 0")
	}
	for k, count := range fakeImp.upsertedBars {
		if count != 1 {
			t.Errorf("expected key %s count to be 1, got %d", k, count)
		}
	}
	fakeImp.mu.Unlock()

	// 2nd tick
	res2, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("second tick: %v", err)
	}
	if len(res2.Items) != 1 || res2.Items[0].Status != "succeeded" {
		t.Fatalf("expected succeeded in 2nd tick, got: %+v", res2.Items)
	}

	fakeImp.mu.Lock()
	keysCount2 := len(fakeImp.upsertedBars)
	if keysCount2 != keysCount1 {
		t.Errorf("expected unique key count to be %d, got %d (duplicated keys detected)", keysCount1, keysCount2)
	}
	for k, count := range fakeImp.upsertedBars {
		if count != 2 {
			t.Errorf("expected key %s count to be 2 after second tick, got %d", k, count)
		}
	}
	fakeImp.mu.Unlock()
}

// --- ExecuteWithResult error propagation tests ---

func TestRunTickResultHandlerErrorDoesNotReexecute(t *testing.T) {
	// When ExecuteWithResult returns (Result{}, false, err), the scheduler
	// must record "failed" WITHOUT calling Execute() again.  This verifies
	// the fix for the original bug where the first error was lost.
	resetActiveJobs()

	wantErrMsg := "import failed: connection reset"
	exec := &fakeExecutor{
		resultFn: func(ctx context.Context, job jobs.Job) (jobs.Result, bool, error) {
			return jobs.Result{}, false, fmt.Errorf(wantErrMsg)
		},
	}

	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "error-test",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"A"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 1,
			},
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec,
	}

	res, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("RunTick: %v", err)
	}

	if len(res.Items) != 1 {
		t.Fatalf("expected 1 item, got %d", len(res.Items))
	}

	item := res.Items[0]
	if item.Status != "failed" {
		t.Errorf("expected status=failed, got %s", item.Status)
	}
	if item.Readiness != "error" {
		t.Errorf("expected readiness=error, got %s", item.Readiness)
	}
	if item.Freshness != "stale" {
		t.Errorf("expected freshness=stale, got %s", item.Freshness)
	}
	if !strings.Contains(item.Error, wantErrMsg) {
		t.Errorf("expected error to contain %q, got %q", wantErrMsg, item.Error)
	}
	if item.ImportSummary != nil {
		t.Errorf("expected nil ImportSummary on failure, got %+v", item.ImportSummary)
	}
}

func TestRunTickResultHandlerSuccessWithBars(t *testing.T) {
	// Verify that a successful ExecuteWithResult returns bars into ImportSummary.
	resetActiveJobs()

	exec := &fakeExecutor{
		resultFn: func(ctx context.Context, job jobs.Job) (jobs.Result, bool, error) {
			return jobs.Result{Bars: 42, GapBars: 2}, true, nil
		},
	}

	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "success-test",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"A"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 1,
			},
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec,
	}

	res, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("RunTick: %v", err)
	}

	if len(res.Items) != 1 {
		t.Fatalf("expected 1 item, got %d", len(res.Items))
	}

	item := res.Items[0]
	if item.Status != "succeeded" {
		t.Errorf("expected status=succeeded, got %s", item.Status)
	}
	if item.ImportSummary == nil {
		t.Fatal("expected ImportSummary, got nil")
	}
	if item.ImportSummary.ImportedBars != 42 {
		t.Errorf("expected ImportedBars=42, got %d", item.ImportSummary.ImportedBars)
	}
	if item.ImportSummary.GapBars != 2 {
		t.Errorf("expected GapBars=2, got %d", item.ImportSummary.GapBars)
	}
}

func TestRunTickNoResultHandlerFallsBackToExecute(t *testing.T) {
	// When resultFn is nil, ExecuteWithResult returns (Result{}, false, nil).
	// The scheduler must fall back to Execute() for success/error.
	fallbackExecuted := false
	exec := &fakeExecutor{
		// resultFn is nil — default legacy path.
		executeFn: func(ctx context.Context, job jobs.Job) error {
			fallbackExecuted = true
			return nil
		},
	}

	cfg := Config{
		Schedules: []ScheduleConfig{
			{
				Name:     "fallback-test",
				Provider: "kis",
				Selector: SelectorConfig{
					Kind:    "watchlist",
					Market:  "KR",
					Venue:   "KRX",
					Name:    "watchlist1",
					Symbols: []string{"A"},
				},
				Timeframe:        "daily",
				Cadence:          "daily",
				Timezone:         "Asia/Seoul",
				BackfillWindow:   "3d",
				ParallelismLimit: 1,
			},
		},
	}

	opts := TickOptions{
		Now:          time.Date(2026, 6, 21, 10, 30, 0, 0, time.UTC),
		Capabilities: CapabilityMap(testCapability()),
		Executor:     exec,
	}

	res, err := RunTick(context.Background(), cfg, opts)
	if err != nil {
		t.Fatalf("RunTick: %v", err)
	}

	if len(res.Items) != 1 {
		t.Fatalf("expected 1 item, got %d", len(res.Items))
	}

	item := res.Items[0]
	if item.Status != "succeeded" {
		t.Errorf("expected status=succeeded, got %s", item.Status)
	}
	if !fallbackExecuted {
		t.Error("expected Execute() to be called as fallback, but it was not")
	}
	if item.ImportSummary != nil {
		t.Errorf("expected nil ImportSummary for no-evidence success, got %+v", item.ImportSummary)
	}
}

func TestRunTickOutputCarriesS02Evidence(t *testing.T) {
	res := TickResult{
		Items: []ItemResult{
			{
				Schedule:  "test-sched",
				Symbol:    "AAPL",
				Status:    "succeeded",
				Provider:  "kis",
				Timeframe: "daily",
				Readiness: "ready",
				Freshness: "fresh",
			},
		},
	}

	var jsonl bytes.Buffer
	if err := WriteTickJSONL(&jsonl, res); err != nil {
		t.Fatalf("WriteTickJSONL: %v", err)
	}

	outStr := jsonl.String()
	for _, expected := range []string{`"status":"succeeded"`, `"readiness":"ready"`, `"freshness":"fresh"`, `"schedule":"test-sched"`, `"symbol":"AAPL"`} {
		if !strings.Contains(outStr, expected) {
			t.Errorf("JSONL output missing expected pattern %q: %s", expected, outStr)
		}
	}

	var text bytes.Buffer
	if err := WriteTickText(&text, res); err != nil {
		t.Fatalf("WriteTickText: %v", err)
	}

	textStr := text.String()
	for _, expected := range []string{"status=succeeded", "readiness=ready", "freshness=fresh", "schedule=test-sched", "symbol=AAPL"} {
		if !strings.Contains(textStr, expected) {
			t.Errorf("Text output missing expected pattern %q: %s", expected, textStr)
		}
	}
}