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

package backtest

import (
	"context"
	"testing"
	"time"

	"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 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)
		}
	}
}