392 lines
9.7 KiB
Go
392 lines
9.7 KiB
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
|
|
}
|
|
|
|
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 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()
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
|