alt/services/worker/internal/backtest/bar_source_test.go

package backtest

import (
	"context"
	"testing"
	"time"

	"git.toki-labs.com/toki/alt/packages/domain/backtest"
	"git.toki-labs.com/toki/alt/packages/domain/market"
)

type sourceStore struct {
	instruments []market.Instrument
	bars        map[market.InstrumentID][]market.Bar
}

func (s *sourceStore) UpsertInstrument(_ context.Context, inst market.Instrument) error {
	s.instruments = append(s.instruments, inst)
	return nil
}

func (s *sourceStore) GetInstrument(_ context.Context, id market.InstrumentID) (market.Instrument, error) {
	for _, inst := range s.instruments {
		if inst.ID == id {
			return inst, nil
		}
	}
	return market.Instrument{}, nil
}

func (s *sourceStore) ListInstruments(_ context.Context) ([]market.Instrument, error) {
	out := make([]market.Instrument, len(s.instruments))
	copy(out, s.instruments)
	return out, nil
}

func (s *sourceStore) UpsertBar(_ context.Context, bar market.Bar) error {
	if s.bars == nil {
		s.bars = make(map[market.InstrumentID][]market.Bar)
	}
	s.bars[bar.InstrumentID] = append(s.bars[bar.InstrumentID], bar)
	return nil
}

func (s *sourceStore) GetBars(_ context.Context, id market.InstrumentID, timeframe market.Timeframe, from, to time.Time) ([]market.Bar, error) {
	var out []market.Bar
	for _, bar := range s.bars[id] {
		if bar.Timeframe != timeframe {
			continue
		}
		if bar.Timestamp.Before(from) || bar.Timestamp.After(to) {
			continue
		}
		out = append(out, bar)
	}
	return out, nil
}

func TestStorageBarSourceReadsMarketBars(t *testing.T) {
	store := &sourceStore{}
	krx := market.Instrument{ID: "KRX:005930", Market: market.MarketKR, Venue: market.VenueKRX, Currency: market.CurrencyKRW}
	nasdaq := market.Instrument{ID: "NASDAQ:AAPL", Market: market.MarketUS, Venue: market.VenueNASDAQ, Currency: market.CurrencyUSD}
	if err := store.UpsertInstrument(context.Background(), krx); err != nil {
		t.Fatalf("upsert krx instrument: %v", err)
	}
	if err := store.UpsertInstrument(context.Background(), nasdaq); err != nil {
		t.Fatalf("upsert us instrument: %v", err)
	}

	day1 := time.Date(2024, 5, 27, 0, 0, 0, 0, time.UTC)
	day2 := time.Date(2024, 5, 28, 0, 0, 0, 0, time.UTC)
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: krx.ID, Timeframe: market.TimeframeDaily, Timestamp: day2}); err != nil {
		t.Fatalf("upsert day2: %v", err)
	}
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: krx.ID, Timeframe: market.TimeframeDaily, Timestamp: day1}); err != nil {
		t.Fatalf("upsert day1: %v", err)
	}
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: nasdaq.ID, Timeframe: market.TimeframeDaily, Timestamp: day1}); err != nil {
		t.Fatalf("upsert us day1: %v", err)
	}

	source := NewStorageBarSource(store, store)
	got, err := source.GetBars(context.Background(), market.MarketKR, market.TimeframeDaily, day1, day2)
	if err != nil {
		t.Fatalf("get bars: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("bars: got %d, want 2", len(got))
	}
	if got[0].Timestamp != day1 || got[1].Timestamp != day2 {
		t.Fatalf("bars are not sorted by timestamp: %+v", got)
	}
	for _, bar := range got {
		if bar.InstrumentID != krx.ID {
			t.Fatalf("unexpected instrument id: %q", bar.InstrumentID)
		}
	}
}

func TestStorageBarSourceKeepsTimeframesSeparate(t *testing.T) {
	store := &sourceStore{}
	inst := market.Instrument{ID: "KRX:005930", Market: market.MarketKR, Venue: market.VenueKRX, Currency: market.CurrencyKRW}
	if err := store.UpsertInstrument(context.Background(), inst); err != nil {
		t.Fatalf("upsert instrument: %v", err)
	}

	ts := time.Date(2024, 5, 27, 0, 0, 0, 0, time.UTC)
	for _, bar := range []market.Bar{
		{InstrumentID: inst.ID, Timeframe: market.TimeframeMonthly, Timestamp: ts},
		{InstrumentID: inst.ID, Timeframe: market.TimeframeDaily, Timestamp: ts},
		{InstrumentID: inst.ID, Timeframe: market.TimeframeMin1, Timestamp: ts},
	} {
		if err := store.UpsertBar(context.Background(), bar); err != nil {
			t.Fatalf("upsert %s bar: %v", bar.Timeframe, err)
		}
	}

	source := NewStorageBarSource(store, store)
	for _, timeframe := range []market.Timeframe{market.TimeframeMonthly, market.TimeframeDaily, market.TimeframeMin1} {
		t.Run(string(timeframe), func(t *testing.T) {
			got, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
				Market:    market.MarketKR,
				Timeframe: timeframe,
				From:      ts,
				To:        ts,
				Selector:  backtest.InputSelector{InstrumentIDs: []market.InstrumentID{inst.ID}},
			})
			if err != nil {
				t.Fatalf("get bars for run: %v", err)
			}
			if len(got) != 1 {
				t.Fatalf("bars: got %d, want 1", len(got))
			}
			if got[0].Timeframe != timeframe {
				t.Fatalf("timeframe: got %q, want %q", got[0].Timeframe, timeframe)
			}
		})
	}
}

func TestStorageBarSourceReturnsOnlyUSBarsForUSMarket(t *testing.T) {
	store := &sourceStore{}
	krx := market.Instrument{ID: "KRX:005930", Market: market.MarketKR, Venue: market.VenueKRX, Currency: market.CurrencyKRW}
	nasdaq := market.Instrument{ID: "NASDAQ:AAPL", Market: market.MarketUS, Venue: market.VenueNASDAQ, Currency: market.CurrencyUSD}
	if err := store.UpsertInstrument(context.Background(), krx); err != nil {
		t.Fatalf("upsert krx instrument: %v", err)
	}
	if err := store.UpsertInstrument(context.Background(), nasdaq); err != nil {
		t.Fatalf("upsert us instrument: %v", err)
	}

	day1 := time.Date(2024, 5, 27, 0, 0, 0, 0, time.UTC)
	day2 := time.Date(2024, 5, 28, 0, 0, 0, 0, time.UTC)
	// Seed KR bars
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: krx.ID, Timeframe: market.TimeframeDaily, Timestamp: day2}); err != nil {
		t.Fatalf("upsert krx day2: %v", err)
	}
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: krx.ID, Timeframe: market.TimeframeDaily, Timestamp: day1}); err != nil {
		t.Fatalf("upsert krx day1: %v", err)
	}
	// Seed US bars
	if err := store.UpsertBar(context.Background(), market.Bar{
		InstrumentID: nasdaq.ID, Timeframe: market.TimeframeDaily,
		Timestamp: day2, Close: market.Price{Currency: market.CurrencyUSD, Amount: market.Decimal{Value: "155"}},
	}); err != nil {
		t.Fatalf("upsert us day2: %v", err)
	}
	if err := store.UpsertBar(context.Background(), market.Bar{
		InstrumentID: nasdaq.ID, Timeframe: market.TimeframeDaily,
		Timestamp: day1, Close: market.Price{Currency: market.CurrencyUSD, Amount: market.Decimal{Value: "150"}},
	}); err != nil {
		t.Fatalf("upsert us day1: %v", err)
	}

	source := NewStorageBarSource(store, store)

	// Query US market — should return only NASDAQ:AAPL bars
	got, err := source.GetBars(context.Background(), market.MarketUS, market.TimeframeDaily, day1, day2)
	if err != nil {
		t.Fatalf("get bars: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("US bars: got %d, want 2", len(got))
	}
	// Verify chronological order
	if got[0].Timestamp != day1 || got[1].Timestamp != day2 {
		t.Fatalf("US bars not sorted by timestamp: %+v", got)
	}
	// All returned bars must be NASDAQ:AAPL
	for _, bar := range got {
		if bar.InstrumentID != nasdaq.ID {
			t.Fatalf("US query returned unexpected instrument id: %q, want %q", bar.InstrumentID, nasdaq.ID)
		}
	}
	// Verify USD currency is preserved in bar Close price
	for _, bar := range got {
		if bar.Close.Currency != market.CurrencyUSD {
			t.Fatalf("US bar Close currency: got %q, want %q", bar.Close.Currency, market.CurrencyUSD)
		}
	}
}

// seedTwoKRInstruments returns a store with two KR instruments, each with one
// daily bar in the same window, so selector tests can prove only the selected
// instrument's bars are returned.
func seedTwoKRInstruments(t *testing.T) (*sourceStore, market.Instrument, market.Instrument, time.Time, time.Time) {
	t.Helper()
	store := &sourceStore{}
	samsung := market.Instrument{
		ID: "KRX:005930", Market: market.MarketKR, Venue: market.VenueKRX, Currency: market.CurrencyKRW,
		Symbol: "005930", ProviderSymbols: map[string]string{"kis": "005930.KS"},
	}
	hynix := market.Instrument{
		ID: "KRX:000660", Market: market.MarketKR, Venue: market.VenueKRX, Currency: market.CurrencyKRW,
		Symbol: "000660", ProviderSymbols: map[string]string{"kis": "000660.KS"},
	}
	for _, inst := range []market.Instrument{samsung, hynix} {
		if err := store.UpsertInstrument(context.Background(), inst); err != nil {
			t.Fatalf("upsert instrument %q: %v", inst.ID, err)
		}
	}
	day1 := time.Date(2024, 5, 27, 0, 0, 0, 0, time.UTC)
	day2 := time.Date(2024, 5, 28, 0, 0, 0, 0, time.UTC)
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: samsung.ID, Timeframe: market.TimeframeDaily, Timestamp: day1}); err != nil {
		t.Fatalf("upsert samsung bar: %v", err)
	}
	if err := store.UpsertBar(context.Background(), market.Bar{InstrumentID: hynix.ID, Timeframe: market.TimeframeDaily, Timestamp: day2}); err != nil {
		t.Fatalf("upsert hynix bar: %v", err)
	}
	return store, samsung, hynix, day1, day2
}

func TestStorageBarSourceSelectsInstrumentIDs(t *testing.T) {
	store, samsung, _, day1, day2 := seedTwoKRInstruments(t)
	source := NewStorageBarSource(store, store)

	got, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{InstrumentIDs: []market.InstrumentID{samsung.ID}},
	})
	if err != nil {
		t.Fatalf("get bars for run: %v", err)
	}
	if len(got) != 1 {
		t.Fatalf("selected bars: got %d, want 1", len(got))
	}
	if got[0].InstrumentID != samsung.ID {
		t.Fatalf("selected instrument id: got %q, want %q", got[0].InstrumentID, samsung.ID)
	}
}

func TestStorageBarSourceSelectsSymbols(t *testing.T) {
	store, _, hynix, day1, day2 := seedTwoKRInstruments(t)
	source := NewStorageBarSource(store, store)

	// Canonical symbol selects the instrument.
	got, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{Symbols: []string{"000660"}},
	})
	if err != nil {
		t.Fatalf("get bars for run by symbol: %v", err)
	}
	if len(got) != 1 || got[0].InstrumentID != hynix.ID {
		t.Fatalf("symbol selection: got %+v, want only %q", got, hynix.ID)
	}

	// Provider symbol value selects the same instrument.
	gotProvider, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{Symbols: []string{"000660.KS"}},
	})
	if err != nil {
		t.Fatalf("get bars for run by provider symbol: %v", err)
	}
	if len(gotProvider) != 1 || gotProvider[0].InstrumentID != hynix.ID {
		t.Fatalf("provider symbol selection: got %+v, want only %q", gotProvider, hynix.ID)
	}
}

func TestStorageBarSourceRejectsMissingSelector(t *testing.T) {
	store, _, _, day1, day2 := seedTwoKRInstruments(t)
	source := NewStorageBarSource(store, store)

	if _, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{InstrumentIDs: []market.InstrumentID{"KRX:999999"}},
	}); err == nil {
		t.Fatal("expected error for unresolved instrument id selector, got nil")
	}

	if _, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{Symbols: []string{"NOPE"}},
	}); err == nil {
		t.Fatal("expected error for unresolved symbol selector, got nil")
	}
}

func TestStorageBarSourceRejectsWhitespaceOnlySelector(t *testing.T) {
	store, _, _, day1, day2 := seedTwoKRInstruments(t)
	source := NewStorageBarSource(store, store)

	if _, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{InstrumentIDs: []market.InstrumentID{" "}},
	}); err == nil {
		t.Fatal("expected error for whitespace-only instrument id, got nil")
	}

	if _, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
		Selector:  backtest.InputSelector{Symbols: []string{"  "}},
	}); err == nil {
		t.Fatal("expected error for whitespace-only symbol, got nil")
	}
}

func TestStorageBarSourceEmptySelectorReadsMarket(t *testing.T) {
	store, _, _, day1, day2 := seedTwoKRInstruments(t)
	source := NewStorageBarSource(store, store)

	got, err := source.GetBarsForRun(context.Background(), backtest.RunSpec{
		Market:    market.MarketKR,
		Timeframe: market.TimeframeDaily,
		From:      day1,
		To:        day2,
	})
	if err != nil {
		t.Fatalf("get bars for run with empty selector: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("empty selector should read the whole market: got %d bars, want 2", len(got))
	}
	if got[0].Timestamp != day1 || got[1].Timestamp != day2 {
		t.Fatalf("bars not sorted by timestamp: %+v", got)
	}
}