From e5dc898fdeb014244a06c8653f420d71fc88e4ac Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Tue, 21 Jul 2026 15:52:05 +0700 Subject: [PATCH] feat(portfolio): track stock and coin cost basis --- cmd/server/main.go | 42 ++++- internal/modules/coin/handlers.go | 14 +- internal/modules/coin/handlers_test.go | 32 +++- internal/modules/coin/portfolio.go | 105 +++++++++++- internal/modules/coin/portfolio_test.go | 16 ++ internal/modules/coin/startup.go | 159 ++++++++++++++++++ internal/modules/coin/startup_mongo_test.go | 60 +++++++ internal/modules/coin/startup_test.go | 141 ++++++++++++++++ internal/modules/coin/views.go | 69 ++++++-- .../modules/coin/views_reply_budget_test.go | 14 ++ internal/modules/gold/portfolio_test.go | 1 + internal/modules/stock/handlers.go | 86 ++++++++-- internal/modules/stock/handlers_test.go | 46 ++++- internal/modules/stock/portfolio.go | 84 ++++++++- internal/modules/stock/portfolio_test.go | 1 + internal/modules/stock/startup.go | 158 +++++++++++++++++ internal/modules/stock/startup_mongo_test.go | 60 +++++++ internal/modules/stock/startup_test.go | 140 +++++++++++++++ internal/modules/stock/stats_test.go | 26 ++- 19 files changed, 1200 insertions(+), 54 deletions(-) create mode 100644 internal/modules/coin/startup.go create mode 100644 internal/modules/coin/startup_mongo_test.go create mode 100644 internal/modules/coin/startup_test.go create mode 100644 internal/modules/stock/startup.go create mode 100644 internal/modules/stock/startup_mongo_test.go create mode 100644 internal/modules/stock/startup_test.go diff --git a/cmd/server/main.go b/cmd/server/main.go index df09e1e..19b2d7f 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -78,15 +78,15 @@ func resolveCommitSHA(envSourceCommit string) string { // import cycle (modules → util → modules). func factories() map[string]modules.Factory { return map[string]modules.Factory{ - "util": util.New, - "misc": misc.New, - "wordle": wordle.New, - "loldle": loldle.New, - lol.CollectionName: lol.New, - "coin": coin.New, - "gold": gold.New, - "stock": stock.New, - "stats": stats.New, + "util": util.New, + "misc": misc.New, + "wordle": wordle.New, + "loldle": loldle.New, + lol.CollectionName: lol.New, + coin.CollectionName: coin.New, + "gold": gold.New, + stock.CollectionName: stock.New, + "stats": stats.New, } } @@ -94,6 +94,7 @@ func factories() map[string]modules.Factory { // shutdown). Atlas SRV DNS + TLS handshake can take a couple seconds on a cold // container; 10s leaves headroom without hiding a wedged cluster. const mongodbInitTimeout = 10 * time.Second +const portfolioMigrationTimeout = 2 * time.Minute func main() { rootCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) @@ -132,6 +133,20 @@ func main() { if err != nil { log.Fatal("module registry build failed", "err", err) } + migrationCtx, cancelMigrations := context.WithTimeout(rootCtx, portfolioMigrationTimeout) + if moduleLoaded(reg, stock.CollectionName) { + if err := stock.InitStore(migrationCtx, provider.Collection(stock.CollectionName), provider.Collection(systemstate.CollectionName), &stock.PriceClient{}); err != nil { + cancelMigrations() + log.Fatal("stock storage init failed", "err", err) + } + } + if moduleLoaded(reg, coin.CollectionName) { + if err := coin.InitStore(migrationCtx, provider.Collection(coin.CollectionName), provider.Collection(systemstate.CollectionName), coin.NewPriceClient()); err != nil { + cancelMigrations() + log.Fatal("coin storage init failed", "err", err) + } + } + cancelMigrations() auth := modules.Auth{BotOwnerID: cfg.BotOwnerID, AdminUserIDs: cfg.AdminUserIDs} modules.Install(b, reg, auth) log.Info("modules loaded", @@ -210,6 +225,15 @@ func main() { } } +func moduleLoaded(reg *modules.Registry, name string) bool { + for _, module := range reg.Modules { + if module.Name == name { + return true + } + } + return false +} + // buildProvider picks the storage backend. Selection order: // 1. Explicit KV_PROVIDER env (memory|mongodb) wins. // 2. Auto-detect: MONGO_URL set → mongodb; otherwise memory. diff --git a/internal/modules/coin/handlers.go b/internal/modules/coin/handlers.go index 756aa63..02ebcf4 100644 --- a/internal/modules/coin/handlers.go +++ b/internal/modules/coin/handlers.go @@ -100,6 +100,9 @@ func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update insufficientBalance = &balance return errInsufficientUSD } + if err := p.AddCostBasis(coin.Symbol, amount); err != nil { + return err + } p.AddAsset(coin.Symbol, qty) return nil }) @@ -149,13 +152,21 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat defer s.locks.Acquire(strconv.FormatInt(userID, 10))() var insufficientHeldQty float64 var insufficientHeld bool + var soldBasis float64 p, err := UpdatePortfolio(ctx, s.store, userID, s.now().UnixMilli(), func(p *Portfolio) error { + heldBefore := p.Assets[coin.Symbol] ok, held := p.DeductAsset(coin.Symbol, qty) if !ok { insufficientHeldQty = held insufficientHeld = true return errInsufficientCoin } + remaining := p.Assets[coin.Symbol] + var basisErr error + soldBasis, basisErr = p.RemoveCostBasis(coin.Symbol, qty, heldBefore, remaining > 0) + if basisErr != nil { + return basisErr + } p.AddUSD(amount) return nil }) @@ -168,7 +179,8 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat } return chathelper.Reply(ctx, b, update.Message, "Sold "+FormatCoinQty(qty)+" "+coin.Symbol+" @ "+FormatUSD(price.USD)+" ("+price.Source+")"+ - "\nReceived: "+FormatUSD(amount)+"\nRemaining: "+FormatUSD(p.USD)) + "\nReceived: "+FormatUSD(amount)+"\nRealized P&L: "+FormatPnLUSD(amount, soldBasis)+ + "\nRemaining: "+FormatUSD(p.USD)) } func formatInsufficientSellMessage(coin CoinSymbol, requestedUSD, heldQty, priceUSD float64) string { diff --git a/internal/modules/coin/handlers_test.go b/internal/modules/coin/handlers_test.go index 33fbcdb..9e71fd0 100644 --- a/internal/modules/coin/handlers_test.go +++ b/internal/modules/coin/handlers_test.go @@ -3,6 +3,7 @@ package coin import ( "context" "errors" + "math" "strings" "testing" "time" @@ -155,13 +156,18 @@ func TestHandleBuyAndSell(t *testing.T) { if p.USD != 500 || p.Assets["BTC"] != 0.01 { t.Fatalf("after buy = %+v", p) } + if p.CostBasis["BTC"] != 500 { + t.Fatalf("buy cost basis = %v, want 500", p.CostBasis["BTC"]) + } + s.prices.(fakePriceFetcher).prices["BTC"] = CoinPrice{USD: 60_000, Source: "Binance"} rb.Reset() - if err := s.handleSell(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/coin_sell 500 BTC")); err != nil { + if err := s.handleSell(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/coin_sell 600 BTC")); err != nil { t.Fatalf("handleSell: %v", err) } rb.AssertSentText(t, "Sold 0.01 BTC") + rb.AssertSentText(t, "Realized P&L: +$100.00 (+20.00%)") p, _ = LoadPortfolio(ctx, s.store, 7, 999) - if p.USD != 1000 || len(p.Assets) != 0 { + if p.USD != 1100 || len(p.Assets) != 0 || len(p.CostBasis) != 0 { t.Fatalf("after sell = %+v", p) } } @@ -339,6 +345,28 @@ func TestStatsWithAndWithoutPrice(t *testing.T) { t.Fatalf("handleStats no price: %v", err) } rb.AssertSentText(t, "price unavailable") + if strings.Contains(rb.LastSent().Text(), "Account P&L: +") || strings.Contains(rb.LastSent().Text(), "Account P&L: -") { + t.Fatalf("partial prices must not show numeric account P&L: %q", rb.LastSent().Text()) + } +} + +func TestStatsTreatsOverflowedValuationAsUnavailable(t *testing.T) { + ctx := context.Background() + s := newTestState(map[string]CoinPrice{"BTC": {USD: math.MaxFloat64, Source: "test"}}, nil) + p := NewPortfolio(1) + p.Assets["BTC"] = 2 + p.CostBasis["BTC"] = 1 + if err := SavePortfolio(ctx, s.store, 7, p); err != nil { + t.Fatal(err) + } + rb := testutil.NewRecordingBot(t) + if err := s.handleStats(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/coin_portfolio")); err != nil { + t.Fatal(err) + } + text := rb.LastSent().Text() + if !strings.Contains(text, "valuation unavailable") || !strings.Contains(text, "Account P&L: unavailable") { + t.Fatalf("overflowed valuation was presented as complete: %q", text) + } } func modDepsForTest() modules.Deps { diff --git a/internal/modules/coin/portfolio.go b/internal/modules/coin/portfolio.go index c2c1e18..d435023 100644 --- a/internal/modules/coin/portfolio.go +++ b/internal/modules/coin/portfolio.go @@ -17,9 +17,10 @@ const portfolioUpdateAttempts = 5 type Store = storage.DocStore[Portfolio] type Portfolio struct { - USD float64 `json:"usd" bson:"usd"` - Assets map[string]float64 `json:"assets" bson:"assets"` - Meta PortfolioMeta `json:"meta" bson:"meta"` + USD float64 `json:"usd" bson:"usd"` + Assets map[string]float64 `json:"assets" bson:"assets"` + CostBasis map[string]float64 `json:"costBasis" bson:"costBasis"` + Meta PortfolioMeta `json:"meta" bson:"meta"` } type PortfolioMeta struct { @@ -28,7 +29,11 @@ type PortfolioMeta struct { } func NewPortfolio(now int64) Portfolio { - return Portfolio{Assets: map[string]float64{}, Meta: PortfolioMeta{CreatedAt: now}} + return Portfolio{ + Assets: map[string]float64{}, + CostBasis: map[string]float64{}, + Meta: PortfolioMeta{CreatedAt: now}, + } } func portfolioKey(userID int64) string { @@ -45,6 +50,9 @@ func LoadPortfolio(ctx context.Context, store Store, userID int64, now int64) (P func SavePortfolio(ctx context.Context, store Store, userID int64, p Portfolio) error { p.normalize() + if err := p.ValidateCostBasis(); err != nil { + return fmt.Errorf("coin: save portfolio %d: %w", userID, err) + } if err := store.Put(ctx, portfolioKey(userID), p); err != nil { return fmt.Errorf("coin: save portfolio %d: %w", userID, err) } @@ -62,6 +70,9 @@ func UpdatePortfolio(ctx context.Context, store Store, userID int64, now int64, return p, err } p.normalize() + if err := p.ValidateCostBasis(); err != nil { + return Portfolio{}, err + } if err := store.PutVersioned(ctx, key, version, p); err == nil { return p, nil } else if !errors.Is(err, storage.ErrConflict) { @@ -75,13 +86,22 @@ func loadPortfolioForUpdate(ctx context.Context, store Store, key string, now in p, version, err := store.Get(ctx, key) switch { case err == nil: + if err := p.validateStoredAssetQuantities(); err != nil { + return Portfolio{}, 0, err + } p.normalize() if p.Assets == nil { p.Assets = map[string]float64{} } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } if p.Meta.CreatedAt == 0 { p.Meta.CreatedAt = now } + if err := p.ValidateCostBasis(); err != nil { + return Portfolio{}, 0, err + } return p, version, nil case errors.Is(err, storage.ErrNotFound): return NewPortfolio(now), 0, nil @@ -90,6 +110,15 @@ func loadPortfolioForUpdate(ctx context.Context, store Store, key string, now in } } +func (p Portfolio) validateStoredAssetQuantities() error { + for symbol, qty := range p.Assets { + if math.IsNaN(qty) || math.IsInf(qty, 0) || qty < 0 { + return fmt.Errorf("coin: %s has invalid quantity", symbol) + } + } + return nil +} + func (p *Portfolio) AddUSD(amount float64) { p.USD += amount p.normalize() @@ -114,6 +143,62 @@ func (p *Portfolio) AddAsset(symbol string, amount float64) { p.normalize() } +func (p *Portfolio) AddCostBasis(symbol string, amount float64) error { + if !isPositiveFinite(amount) { + return fmt.Errorf("coin: invalid purchase cost basis") + } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + next := p.CostBasis[symbol] + amount + if !isPositiveFinite(next) { + return fmt.Errorf("coin: cost basis overflows") + } + p.CostBasis[symbol] = next + return nil +} + +func (p *Portfolio) RemoveCostBasis(symbol string, sold, held float64, holdingRemains bool) (float64, error) { + if !isPositiveFinite(sold) || !isPositiveFinite(held) || sold > held+coinDustEpsilon { + return 0, fmt.Errorf("coin: invalid cost basis quantities") + } + basis := p.CostBasis[symbol] + if !isPositiveFinite(basis) { + return 0, fmt.Errorf("coin: missing cost basis for %s", symbol) + } + if !holdingRemains { + delete(p.CostBasis, symbol) + return basis, nil + } + removed := basis * (sold / held) + remaining := basis - removed + if !isPositiveFinite(removed) || !isPositiveFinite(remaining) { + return 0, fmt.Errorf("coin: invalid remaining cost basis") + } + p.CostBasis[symbol] = remaining + return removed, nil +} + +func (p Portfolio) ValidateCostBasis() error { + for symbol, basis := range p.CostBasis { + if !isPositiveFinite(basis) { + return fmt.Errorf("coin: %s has invalid cost basis", symbol) + } + if p.Assets[symbol] <= 0 { + return fmt.Errorf("coin: %s has cost basis without a holding", symbol) + } + } + for symbol, qty := range p.Assets { + if qty <= 0 { + continue + } + if !isPositiveFinite(p.CostBasis[symbol]) { + return fmt.Errorf("coin: holding %s has missing or invalid cost basis", symbol) + } + } + return nil +} + func (p *Portfolio) DeductAsset(symbol string, amount float64) (ok bool, held float64) { if p.Assets == nil { p.Assets = map[string]float64{} @@ -142,6 +227,18 @@ func (p *Portfolio) normalize() { p.Assets[symbol] = amount } } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + for symbol, basis := range p.CostBasis { + if basis == 0 { + continue + } + if math.IsNaN(basis) || math.IsInf(basis, 0) || basis < 0 { + continue + } + p.CostBasis[symbol] = basis + } } func normalizeAmount(n float64) float64 { diff --git a/internal/modules/coin/portfolio_test.go b/internal/modules/coin/portfolio_test.go index cccb389..206e4ba 100644 --- a/internal/modules/coin/portfolio_test.go +++ b/internal/modules/coin/portfolio_test.go @@ -67,6 +67,22 @@ func TestNormalizeAmountSpecialValues(t *testing.T) { } } +func TestLoadPortfolioRejectsCorruptStoredQuantityBeforeNormalization(t *testing.T) { + ctx := context.Background() + store := &corruptQuantityStore{Store: newCoinStore()} + if _, err := LoadPortfolio(ctx, store, 7, 1); err == nil { + t.Fatal("LoadPortfolio silently normalized a corrupt quantity") + } +} + +type corruptQuantityStore struct{ Store } + +func (s *corruptQuantityStore) Get(context.Context, string) (Portfolio, int64, error) { + p := NewPortfolio(1) + p.Assets["BTC"] = math.NaN() + return p, 1, nil +} + // conflictOnceStore wraps a real typed store and forces exactly one write // conflict (after committing a competing value) before delegating, to exercise // UpdatePortfolio's optimistic-lock retry. diff --git a/internal/modules/coin/startup.go b/internal/modules/coin/startup.go new file mode 100644 index 0000000..b485f04 --- /dev/null +++ b/internal/modules/coin/startup.go @@ -0,0 +1,159 @@ +package coin + +import ( + "context" + "errors" + "fmt" + "sort" + "time" + + "github.com/tiennm99/miti99bot/internal/log" + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +const ( + CollectionName = "coin" + costBasisMigrationKey = "migration:coin-cost-basis-v1" + costBasisMigrationRetries = 5 +) + +type MigrationPriceFetcher interface { + FetchUSD(context.Context, CoinSymbol) (CoinPrice, error) +} + +// InitStore seeds legacy coin holdings at current market prices and verifies +// the invariant again on every startup, even after the audit marker exists. +func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection, prices MigrationPriceFetcher) error { + docs := storage.Typed[Portfolio](portfolioColl) + system := systemstate.New(systemColl) + marker, markerExists, err := system.Get(ctx, costBasisMigrationKey) + if err != nil { + return fmt.Errorf("coin cost basis migration: read marker: %w", err) + } + keys, err := docs.List(ctx, "user:") + if err != nil { + return fmt.Errorf("coin cost basis migration: list portfolios: %w", err) + } + missing := map[string]bool{} + for _, key := range keys { + p, _, err := docs.Get(ctx, key) + if err != nil { + return fmt.Errorf("coin cost basis migration: read %s: %w", key, err) + } + if err := inspectLegacyPortfolio(p, missing); err != nil { + return fmt.Errorf("coin cost basis migration: %s: %w", key, err) + } + } + + quotes := map[string]float64{} + for _, symbol := range sortedMissingSymbols(missing) { + coin, err := ResolveCoinSymbol(symbol) + if err != nil { + return fmt.Errorf("coin cost basis migration: resolve %s: %w", symbol, err) + } + price, err := prices.FetchUSD(ctx, coin) + if err != nil { + return fmt.Errorf("coin cost basis migration: fetch %s quote: %w", symbol, err) + } + if !isPositiveFinite(price.USD) { + return fmt.Errorf("coin cost basis migration: no valid quote for %s", symbol) + } + quotes[symbol] = price.USD + } + + var migrated int64 + for index, key := range keys { + count, err := migrateLegacyPortfolio(ctx, docs, key, quotes) + if err != nil { + return err + } + migrated += int64(count) + if count > 0 { + log.Info("coin cost basis migrated", "portfolio", index+1, "total", len(keys), "positions", count) + } + } + now := time.Now().UnixMilli() + if markerExists && marker.Status == "completed" && migrated == 0 { + return nil + } + if !markerExists { + marker = systemstate.Record{Kind: "migration", Name: "coin cost basis v1", CompletedAt: now} + } + marker.Status = "completed" + marker.Count += migrated + marker.UpdatedAt = now + if err := system.Put(ctx, costBasisMigrationKey, marker); err != nil { + return fmt.Errorf("coin cost basis migration: write marker: %w", err) + } + return nil +} + +func inspectLegacyPortfolio(p Portfolio, missing map[string]bool) error { + for symbol, basis := range p.CostBasis { + if !isPositiveFinite(basis) { + return fmt.Errorf("%s has invalid cost basis", symbol) + } + if p.Assets[symbol] <= 0 { + return fmt.Errorf("%s has cost basis without a holding", symbol) + } + } + for symbol, qty := range p.Assets { + if !isPositiveFinite(qty) { + if qty == 0 { + continue + } + return fmt.Errorf("%s has invalid quantity", symbol) + } + coin, err := ResolveCoinSymbol(symbol) + if err != nil || coin.Symbol != symbol { + return fmt.Errorf("%q is not a canonical coin symbol", symbol) + } + if _, ok := p.CostBasis[symbol]; !ok { + missing[symbol] = true + } + } + return nil +} + +func migrateLegacyPortfolio(ctx context.Context, docs storage.DocStore[Portfolio], key string, quotes map[string]float64) (int, error) { + for attempt := 0; attempt < costBasisMigrationRetries; attempt++ { + p, version, err := docs.Get(ctx, key) + if err != nil { + return 0, fmt.Errorf("coin cost basis migration: read %s: %w", key, err) + } + missing := map[string]bool{} + if err := inspectLegacyPortfolio(p, missing); err != nil { + return 0, fmt.Errorf("coin cost basis migration: %s: %w", key, err) + } + if len(missing) == 0 { + return 0, nil + } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + for symbol := range missing { + quote := quotes[symbol] + basis := p.Assets[symbol] * quote + if !isPositiveFinite(quote) || !isPositiveFinite(basis) { + return 0, fmt.Errorf("coin cost basis migration: no cached valid quote for %s", symbol) + } + p.CostBasis[symbol] = basis + } + if err := docs.PutVersioned(ctx, key, version, p); err == nil { + return len(missing), nil + } else if !errors.Is(err, storage.ErrConflict) { + return 0, fmt.Errorf("coin cost basis migration: write %s: %w", key, err) + } + } + return 0, fmt.Errorf("coin cost basis migration: write %s: %w", key, storage.ErrConflict) +} + +func sortedMissingSymbols(missing map[string]bool) []string { + symbols := make([]string, 0, len(missing)) + for symbol := range missing { + symbols = append(symbols, symbol) + } + sort.Strings(symbols) + return symbols +} diff --git a/internal/modules/coin/startup_mongo_test.go b/internal/modules/coin/startup_mongo_test.go new file mode 100644 index 0000000..9481ba9 --- /dev/null +++ b/internal/modules/coin/startup_mongo_test.go @@ -0,0 +1,60 @@ +package coin + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" + "github.com/tiennm99/miti99bot/internal/testutil/mongotest" +) + +var mongoTests mongotest.Manager + +func TestMain(m *testing.M) { + os.Exit(mongoTests.Run(m)) +} + +func TestInitStoreMigratesCoinBasisInMongoDB(t *testing.T) { + ctx, portfolioColl, systemColl := setupMongoCoinTest(t) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["BTC"] = 0.25 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + prices := &migrationCoinPrices{quotes: map[string]float64{"BTC": 100_000}} + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore: %v", err) + } + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore second run: %v", err) + } + got, _, err := docs.Get(ctx, "user:7") + if err != nil || got.CostBasis["BTC"] != 25_000 { + t.Fatalf("portfolio=%+v err=%v", got, err) + } +} + +func setupMongoCoinTest(t *testing.T) (context.Context, storage.Collection, storage.Collection) { + t.Helper() + uri := mongoTests.URI(t) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + client, err := storage.NewMongoClient(ctx, uri) + if err != nil { + t.Fatal(err) + } + db := client.Database(fmt.Sprintf("miti99bot_coin_test_%d", time.Now().UnixNano())) + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + _ = db.Drop(cleanupCtx) + _ = client.Disconnect(cleanupCtx) + }) + provider := storage.NewMongoProvider(db) + return ctx, provider.Collection(CollectionName), provider.Collection(systemstate.CollectionName) +} diff --git a/internal/modules/coin/startup_test.go b/internal/modules/coin/startup_test.go new file mode 100644 index 0000000..e1e9bba --- /dev/null +++ b/internal/modules/coin/startup_test.go @@ -0,0 +1,141 @@ +package coin + +import ( + "context" + "errors" + "math" + "testing" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +type migrationCoinPrices struct { + quotes map[string]float64 + err error + calls int +} + +func (f *migrationCoinPrices) FetchUSD(_ context.Context, coin CoinSymbol) (CoinPrice, error) { + f.calls++ + if f.err != nil { + return CoinPrice{}, f.err + } + return CoinPrice{Symbol: coin.Symbol, USD: f.quotes[coin.Symbol], Source: "test"}, nil +} + +func TestInitStoreMigratesLegacyCoinBasisIdempotently(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + portfolioColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["BTC"] = 0.25 + legacy.Assets["ETH"] = 2 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + prices := &migrationCoinPrices{quotes: map[string]float64{"BTC": 100_000, "ETH": 3_000}} + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore: %v", err) + } + got, _, err := docs.Get(ctx, "user:7") + if err != nil { + t.Fatal(err) + } + if got.CostBasis["BTC"] != 25_000 || got.CostBasis["ETH"] != 6_000 { + t.Fatalf("CostBasis = %#v", got.CostBasis) + } + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore second run: %v", err) + } + if prices.calls != 2 { + t.Fatalf("quote calls = %d, want one per symbol on first run", prices.calls) + } +} + +func TestInitStoreCoinRequiresCompleteQuotesBeforeWriting(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + portfolioColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["BTC"] = 1 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + if err := InitStore(ctx, portfolioColl, systemColl, &migrationCoinPrices{quotes: map[string]float64{}}); err == nil { + t.Fatal("InitStore succeeded without a valid quote") + } + got, _, _ := docs.Get(ctx, "user:7") + if len(got.CostBasis) != 0 { + t.Fatalf("partial migration wrote basis: %#v", got.CostBasis) + } + if _, _, err := storage.Typed[systemstate.Record](systemColl).Get(ctx, costBasisMigrationKey); !errors.Is(err, storage.ErrNotFound) { + t.Fatalf("marker err = %v, want ErrNotFound", err) + } +} + +func TestCoinWeightedAverageBasisAndDustExit(t *testing.T) { + p := NewPortfolio(1) + p.AddAsset("BTC", 0.3) + if err := p.AddCostBasis("BTC", 12_000); err != nil { + t.Fatal(err) + } + removed, err := p.RemoveCostBasis("BTC", 0.1, 0.3, true) + if err != nil { + t.Fatal(err) + } + if math.Abs(removed-4_000) > 1e-9 || math.Abs(p.CostBasis["BTC"]-8_000) > 1e-9 { + t.Fatalf("removed=%v remaining=%v", removed, p.CostBasis["BTC"]) + } + removed, err = p.RemoveCostBasis("BTC", 0.2, 0.2, false) + if err != nil || math.Abs(removed-8_000) > 1e-9 { + t.Fatalf("full exit removed=%v err=%v", removed, err) + } + if _, exists := p.CostBasis["BTC"]; exists { + t.Fatal("full exit retained basis") + } +} + +func TestCoinMigrationScansRowsAddedAfterCompletionMarker(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + portfolioColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + prices := &migrationCoinPrices{quotes: map[string]float64{"BTC": 100_000}} + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatal(err) + } + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["BTC"] = 0.5 + if err := docs.Put(ctx, "user:8", legacy); err != nil { + t.Fatal(err) + } + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatal(err) + } + got, _, err := docs.Get(ctx, "user:8") + if err != nil || got.CostBasis["BTC"] != 50_000 { + t.Fatalf("portfolio=%+v err=%v", got, err) + } +} + +func TestCoinMigrationHonorsCancelledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + provider := storage.NewMemoryProvider() + docs := storage.Typed[Portfolio](provider.Collection(CollectionName)) + legacy := NewPortfolio(1) + legacy.Assets["BTC"] = 1 + if err := docs.Put(context.Background(), "user:7", legacy); err != nil { + t.Fatal(err) + } + prices := &migrationCoinPrices{err: context.Canceled} + if err := InitStore(ctx, provider.Collection(CollectionName), provider.Collection(systemstate.CollectionName), prices); !errors.Is(err, context.Canceled) { + t.Fatalf("InitStore err=%v, want context.Canceled", err) + } +} diff --git a/internal/modules/coin/views.go b/internal/modules/coin/views.go index e26b4b7..cbc1bfc 100644 --- a/internal/modules/coin/views.go +++ b/internal/modules/coin/views.go @@ -3,6 +3,7 @@ package coin import ( "context" "sort" + "strconv" "strings" "github.com/go-telegram/bot" @@ -22,8 +23,11 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda log.Error("coin_load_portfolio", "user", userID, "err", err) return chathelper.Reply(ctx, b, update.Message, "Could not load coin portfolio. Try again later.") } - lines := []string{"Coin Account Summary\n", "USD: " + FormatUSD(p.USD)} + header := []string{"Coin Account Summary", "USD: " + FormatUSD(p.USD)} + var positions []string totalValue := p.USD + totalBasis := 0.0 + missingPrice := false // Fetch sequentially (not concurrently) so the price client's keep-alive // connection pool is reused across coins rather than opening N simultaneous @@ -34,23 +38,49 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda defer cancel() for _, symbol := range sortedAssetSymbols(p.Assets) { held := p.Assets[symbol] - line := symbol + ": " + FormatCoinQty(held) + basis := p.CostBasis[symbol] + average := basis / held + line := symbol + ": " + FormatCoinQty(held) + " | Avg " + FormatUSD(average) if coin, err := ResolveCoinSymbol(symbol); err == nil { - if price, err := s.prices.FetchUSD(fetchCtx, coin); err == nil { + if price, err := s.prices.FetchUSD(fetchCtx, coin); err == nil && isPositiveFinite(price.USD) { value := held * price.USD + if !isPositiveFinite(value) || !isPositiveFinite(average) { + missingPrice = true + positions = append(positions, symbol+": "+FormatCoinQty(held)+" | valuation unavailable") + continue + } totalValue += value - line += " = " + FormatUSD(value) + " @ " + FormatUSD(price.USD) + " (" + price.Source + ")" + totalBasis += basis + line += " | Now " + FormatUSD(price.USD) + " (" + price.Source + ")" + + " | Value " + FormatUSD(value) + " | Unrealized P&L " + FormatPnLUSD(value, basis) } else { log.Error("coin_fetch_price", "symbol", symbol, "err", err) - line += " (price unavailable)" + missingPrice = true + line += " | price unavailable" } + } else { + missingPrice = true + line += " | price unavailable" } - lines = append(lines, line) + positions = append(positions, line) } - lines = append(lines, "Total value: "+FormatUSD(totalValue)) - lines = append(lines, "Invested: "+FormatUSD(p.Meta.Invested)) - lines = append(lines, "P&L: "+FormatPnLUSD(totalValue, p.Meta.Invested)) - return chathelper.Reply(ctx, b, update.Message, strings.Join(lines, "\n")) + var summary []string + if missingPrice { + summary = []string{ + "Priced value (partial): " + FormatUSD(totalValue), + "Unrealized P&L (priced positions): " + FormatPnLUSD(totalValue-p.USD, totalBasis), + "Invested: " + FormatUSD(p.Meta.Invested), + "Account P&L: unavailable until all positions have prices", + } + } else { + summary = []string{ + "Total value: " + FormatUSD(totalValue), + "Invested: " + FormatUSD(p.Meta.Invested), + "Unrealized P&L: " + FormatPnLUSD(totalValue-p.USD, totalBasis), + "Account P&L: " + FormatPnLUSD(totalValue, p.Meta.Invested), + } + } + return chathelper.Reply(ctx, b, update.Message, boundedPortfolioReply(header, positions, summary)) } func sortedAssetSymbols(assets map[string]float64) []string { @@ -63,3 +93,22 @@ func sortedAssetSymbols(assets map[string]float64) []string { sort.Strings(symbols) return symbols } + +const portfolioReplyLimit = 4000 + +func boundedPortfolioReply(header, positions, summary []string) string { + omitted := 0 + for { + lines := append(append(append([]string{}, header...), positions...), summary...) + if omitted > 0 { + insertAt := len(header) + len(positions) + lines = append(lines[:insertAt], append([]string{"… " + strconv.Itoa(omitted) + " position(s) omitted"}, lines[insertAt:]...)...) + } + reply := strings.Join(lines, "\n") + if len(reply) <= portfolioReplyLimit || len(positions) == 0 { + return reply + } + positions = positions[:len(positions)-1] + omitted++ + } +} diff --git a/internal/modules/coin/views_reply_budget_test.go b/internal/modules/coin/views_reply_budget_test.go index d274493..c192d9e 100644 --- a/internal/modules/coin/views_reply_budget_test.go +++ b/internal/modules/coin/views_reply_budget_test.go @@ -60,3 +60,17 @@ func TestHandleStatsDeliversReplyWhenUpstreamHangs(t *testing.T) { t.Fatalf("reply missing summary / degraded line; got:\n%s", sent) } } + +func TestCoinPortfolioReplyStaysWithinTelegramBudget(t *testing.T) { + positions := make([]string, 200) + for i := range positions { + positions[i] = strings.Repeat("position-data-", 20) + } + reply := boundedPortfolioReply([]string{"header"}, positions, []string{"summary"}) + if len(reply) > portfolioReplyLimit { + t.Fatalf("reply length = %d, limit = %d", len(reply), portfolioReplyLimit) + } + if !strings.Contains(reply, "position(s) omitted") || !strings.Contains(reply, "summary") { + t.Fatalf("bounded reply lost omission marker or summary: %q", reply) + } +} diff --git a/internal/modules/gold/portfolio_test.go b/internal/modules/gold/portfolio_test.go index 1d9466b..8ee1fc9 100644 --- a/internal/modules/gold/portfolio_test.go +++ b/internal/modules/gold/portfolio_test.go @@ -235,6 +235,7 @@ func TestStockAndGoldPortfolioKeysDoNotCollide(t *testing.T) { stockStore := storage.Typed[stockmod.Portfolio](provider.Collection("stock")) stockPortfolio := stockmod.NewPortfolio(1) stockPortfolio.AddAsset("TCB", 100) + stockPortfolio.CostBasis["TCB"] = 3_000_000 if err := stockmod.SavePortfolio(ctx, stockStore, 7, stockPortfolio); err != nil { t.Fatalf("save stock: %v", err) } diff --git a/internal/modules/stock/handlers.go b/internal/modules/stock/handlers.go index a0299cd..a115a7a 100644 --- a/internal/modules/stock/handlers.go +++ b/internal/modules/stock/handlers.go @@ -4,6 +4,7 @@ import ( "context" "errors" "math" + "sort" "strconv" "strings" "time" @@ -179,6 +180,10 @@ func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update return chathelper.Reply(ctx, b, update.Message, "Insufficient VND. Need "+FormatVND(cost)+", have "+FormatVND(balance)+".") } + if err := p.AddCostBasis(symbol, cost); err != nil { + log.Error("stock_add_cost_basis", "user", userID, "ticker", symbol, "err", err) + return chathelper.Reply(ctx, b, update.Message, "Could not record purchase cost. Try again later.") + } p.AddAsset(symbol, qty) if err := SavePortfolio(ctx, s.store, userID, p); err != nil { log.Error("stock_save_portfolio", "user", userID, "err", err) @@ -232,12 +237,18 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat log.Error("stock_load_portfolio", "user", userID, "err", err) return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.") } + heldBefore := p.Assets[symbol] ok, held := p.DeductAsset(symbol, qty) if !ok { return chathelper.Reply(ctx, b, update.Message, "Insufficient "+symbol+". You have: "+FormatStock(float64(held))) } revenue := float64(qty) * price + soldBasis, err := p.RemoveCostBasis(symbol, qty, heldBefore) + if err != nil { + log.Error("stock_remove_cost_basis", "user", userID, "ticker", symbol, "err", err) + return chathelper.Reply(ctx, b, update.Message, "Portfolio cost basis is unavailable. Restart the bot or contact the owner.") + } p.AddCurrency("VND", revenue) if err := SavePortfolio(ctx, s.store, userID, p); err != nil { log.Error("stock_save_portfolio", "user", userID, "err", err) @@ -246,6 +257,7 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat return chathelper.Reply(ctx, b, update.Message, "Sold "+FormatStock(float64(qty))+" "+symbol+ " @ "+FormatVND(price)+"\nRevenue: "+FormatVND(revenue)+ + "\nRealized P&L: "+FormatPnL(revenue, soldBasis)+ "\nRemaining: "+FormatVND(p.Currency["VND"])) } @@ -450,14 +462,12 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.") } - var lines []string - lines = append(lines, "📊 Portfolio Summary\n") + header := []string{"📊 Portfolio Summary", "VND: " + FormatVND(p.Currency["VND"])} + var positions []string totalValue := 0.0 - - if vnd := p.Currency["VND"]; vnd > 0 { - totalValue += vnd - lines = append(lines, "VND: "+FormatVND(vnd)) - } + totalValue += p.Currency["VND"] + totalBasis := 0.0 + missingPrice := false // Filter out zero-balance assets (DeductAsset removes them, but defensive). type held struct { @@ -466,13 +476,14 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda } var heldList []held for sym, qty := range p.Assets { - if qty != 0 { + if qty > 0 { heldList = append(heldList, held{sym, qty}) } } + sort.Slice(heldList, func(i, j int) bool { return heldList[i].symbol < heldList[j].symbol }) if len(heldList) > 0 { - lines = append(lines, "\nStocks:") + header = append(header, "", "Stocks:") fetchCtx, cancel := chathelper.FetchContext(ctx) defer cancel() @@ -487,18 +498,59 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda for _, h := range heldList { price := prices[h.symbol] - if fetchErr != nil || price <= 0 { - lines = append(lines, " "+h.symbol+" x"+FormatStock(float64(h.qty))+" (no price)") + basis := p.CostBasis[h.symbol] + average := basis / float64(h.qty) + if !isPositiveFiniteCost(price) { + missingPrice = true + positions = append(positions, " "+h.symbol+" x"+FormatStock(float64(h.qty))+ + " | Avg "+FormatVND(average)+" | price unavailable") continue } val := float64(h.qty) * price + if !isPositiveFiniteCost(val) || !isPositiveFiniteCost(average) { + missingPrice = true + positions = append(positions, " "+h.symbol+" x"+FormatStock(float64(h.qty))+" | valuation unavailable") + continue + } totalValue += val - lines = append(lines, " "+h.symbol+" x"+FormatStock(float64(h.qty))+ - " @ "+FormatVND(price)+" = "+FormatVND(val)) + totalBasis += basis + positions = append(positions, " "+h.symbol+" x"+FormatStock(float64(h.qty))+ + " | Avg "+FormatVND(average)+" | Now "+FormatVND(price)+ + " | Value "+FormatVND(val)+" | Unrealized P&L "+FormatPnL(val, basis)) } } - lines = append(lines, "\nTotal value: "+FormatVND(totalValue)) - lines = append(lines, "Invested: "+FormatVND(p.Meta.Invested)) - lines = append(lines, "P&L: "+FormatPnL(totalValue, p.Meta.Invested)) - return chathelper.Reply(ctx, b, update.Message, strings.Join(lines, "\n")) + var summary []string + if missingPrice { + summary = append(summary, + "", "Priced value (partial): "+FormatVND(totalValue), + "Unrealized P&L (priced positions): "+FormatPnL(totalValue-p.Currency["VND"], totalBasis), + "Invested: "+FormatVND(p.Meta.Invested), + "Account P&L: unavailable until all positions have prices") + } else { + summary = append(summary, + "", "Total value: "+FormatVND(totalValue), + "Invested: "+FormatVND(p.Meta.Invested), + "Unrealized P&L: "+FormatPnL(totalValue-p.Currency["VND"], totalBasis), + "Account P&L: "+FormatPnL(totalValue, p.Meta.Invested)) + } + return chathelper.Reply(ctx, b, update.Message, boundedPortfolioReply(header, positions, summary)) +} + +const portfolioReplyLimit = 4000 + +func boundedPortfolioReply(header, positions, summary []string) string { + omitted := 0 + for { + lines := append(append(append([]string{}, header...), positions...), summary...) + if omitted > 0 { + insertAt := len(header) + len(positions) + lines = append(lines[:insertAt], append([]string{" … " + strconv.Itoa(omitted) + " position(s) omitted"}, lines[insertAt:]...)...) + } + reply := strings.Join(lines, "\n") + if len(reply) <= portfolioReplyLimit || len(positions) == 0 { + return reply + } + positions = positions[:len(positions)-1] + omitted++ + } } diff --git a/internal/modules/stock/handlers_test.go b/internal/modules/stock/handlers_test.go index 59b19a1..67afb5f 100644 --- a/internal/modules/stock/handlers_test.go +++ b/internal/modules/stock/handlers_test.go @@ -3,6 +3,7 @@ package stock import ( "context" "errors" + "fmt" "net/http" "net/http/httptest" "strings" @@ -75,6 +76,42 @@ func TestHandlePriceUsage(t *testing.T) { rb.AssertSentText(t, "Usage: /stock_price ") } +func TestHandleBuyAndPartialSellTracksCostBasisAndRealizedPnL(t *testing.T) { + ctx := context.Background() + price := 30_000.0 + priceSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"data":{"matchedPrice":%v}}`, price) + })) + t.Cleanup(priceSrv.Close) + s := &state{ + store: newStockStore(), + prices: &PriceClient{HTTP: priceSrv.Client(), URL: priceSrv.URL}, + nowFn: func() time.Time { return time.UnixMilli(123) }, + } + rb := testutil.NewRecordingBot(t) + if err := s.handleTopup(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/stock_topup 10000000")); err != nil { + t.Fatal(err) + } + if err := s.handleBuy(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/stock_buy 100 TCB")); err != nil { + t.Fatal(err) + } + p, err := LoadPortfolio(ctx, s.store, 7, 999) + if err != nil || p.CostBasis["TCB"] != 3_000_000 { + t.Fatalf("after buy=%+v err=%v", p, err) + } + price = 36_000 + rb.Reset() + if err := s.handleSell(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/stock_sell 40 TCB")); err != nil { + t.Fatal(err) + } + rb.AssertSentText(t, "Realized P&L: +240.000 VND (+20.00%)") + p, err = LoadPortfolio(ctx, s.store, 7, 999) + if err != nil || p.Assets["TCB"] != 60 || p.CostBasis["TCB"] != 1_800_000 { + t.Fatalf("after sell=%+v err=%v", p, err) + } +} + func TestMutableHandlersRejectExtraArgs(t *testing.T) { ctx := context.Background() s := &state{ @@ -207,6 +244,7 @@ func seedStockPortfolio(t *testing.T, store Store, userID int64, held int64, bal t.Helper() p := NewPortfolio(123) p.Assets["TCB"] = held + p.CostBasis["TCB"] = float64(held) * 30_000 p.Currency["VND"] = balance if err := SavePortfolio(context.Background(), store, userID, p); err != nil { t.Fatalf("seed portfolio: %v", err) @@ -232,6 +270,9 @@ func TestHandleCashDividendAllowsRepeatedManualAdjustments(t *testing.T) { if got, want := p.Currency["VND"], float64(418000); got != want { t.Fatalf("balance = %v, want %v", got, want) } + if p.CostBasis["TCB"] != 139*30_000 { + t.Fatalf("cash dividend changed cost basis: %v", p.CostBasis["TCB"]) + } } func TestHandleCashDividendRejectsInexactBalanceSum(t *testing.T) { @@ -269,6 +310,9 @@ func TestHandleShareDividendPreservesRatioAndFloors(t *testing.T) { if got, want := p.Assets["TCB"], int64(152); got != want { t.Fatalf("holding = %d, want %d", got, want) } + if p.CostBasis["TCB"] != 139*30_000 { + t.Fatalf("share dividend changed total cost basis: %v", p.CostBasis["TCB"]) + } rb.AssertSentText(t, "Share dividend (100:10): +13 TCB") } @@ -335,7 +379,7 @@ func TestHandleCombinedDividendUsesPreEventHoldingAndOneSave(t *testing.T) { t.Fatalf("store writes = %d, want 1", store.puts) } p, _ := LoadPortfolio(ctx, base, 7, 999) - if p.Assets["TCB"] != 152 || p.Currency["VND"] != 209500 { + if p.Assets["TCB"] != 152 || p.Currency["VND"] != 209500 || p.CostBasis["TCB"] != 139*30_000 { t.Fatalf("portfolio = %+v", p) } rb.AssertSentText(t, "Dividend for TCB (100:10)") diff --git a/internal/modules/stock/portfolio.go b/internal/modules/stock/portfolio.go index 79f7061..68c4160 100644 --- a/internal/modules/stock/portfolio.go +++ b/internal/modules/stock/portfolio.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "math" "strconv" "github.com/tiennm99/miti99bot/internal/storage" @@ -15,9 +16,10 @@ type Store = storage.DocStore[Portfolio] // Portfolio is the per-user stock state. Currency is a map for forward- // compat with USD/EUR (currently VND-only). Assets is a flat ticker→qty map. type Portfolio struct { - Currency map[string]float64 `json:"currency" bson:"currency"` - Assets map[string]int64 `json:"assets" bson:"assets"` - Meta PortfolioMeta `json:"meta" bson:"meta"` + Currency map[string]float64 `json:"currency" bson:"currency"` + Assets map[string]int64 `json:"assets" bson:"assets"` + CostBasis map[string]float64 `json:"costBasis" bson:"costBasis"` + Meta PortfolioMeta `json:"meta" bson:"meta"` } // PortfolioMeta tracks invested cost basis for P&L. CreatedAt is purely @@ -32,9 +34,10 @@ type PortfolioMeta struct { // nil-map panics. func NewPortfolio(now int64) Portfolio { return Portfolio{ - Currency: map[string]float64{"VND": 0}, - Assets: map[string]int64{}, - Meta: PortfolioMeta{Invested: 0, CreatedAt: now}, + Currency: map[string]float64{"VND": 0}, + Assets: map[string]int64{}, + CostBasis: map[string]float64{}, + Meta: PortfolioMeta{Invested: 0, CreatedAt: now}, } } @@ -57,6 +60,12 @@ func LoadPortfolio(ctx context.Context, store Store, userID int64, now int64) (P if p.Assets == nil { p.Assets = map[string]int64{} } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + if err := p.ValidateCostBasis(); err != nil { + return Portfolio{}, err + } return p, nil case errors.Is(err, storage.ErrNotFound): return NewPortfolio(now), nil @@ -65,8 +74,71 @@ func LoadPortfolio(ctx context.Context, store Store, userID int64, now int64) (P } } +func (p *Portfolio) AddCostBasis(symbol string, amount float64) error { + if !isPositiveFiniteCost(amount) { + return fmt.Errorf("stock: invalid purchase cost basis") + } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + next := p.CostBasis[symbol] + amount + if !isPositiveFiniteCost(next) { + return fmt.Errorf("stock: cost basis overflows") + } + p.CostBasis[symbol] = next + return nil +} + +func (p *Portfolio) RemoveCostBasis(symbol string, sold, held int64) (float64, error) { + if sold <= 0 || held <= 0 || sold > held { + return 0, fmt.Errorf("stock: invalid cost basis quantities") + } + basis := p.CostBasis[symbol] + if !isPositiveFiniteCost(basis) { + return 0, fmt.Errorf("stock: missing cost basis for %s", symbol) + } + if sold == held { + delete(p.CostBasis, symbol) + return basis, nil + } + removed := basis * (float64(sold) / float64(held)) + remaining := basis - removed + if !isPositiveFiniteCost(removed) || !isPositiveFiniteCost(remaining) { + return 0, fmt.Errorf("stock: invalid remaining cost basis") + } + p.CostBasis[symbol] = remaining + return removed, nil +} + +func (p Portfolio) ValidateCostBasis() error { + for symbol, basis := range p.CostBasis { + if !isPositiveFiniteCost(basis) { + return fmt.Errorf("stock: %s has invalid cost basis", symbol) + } + if p.Assets[symbol] <= 0 { + return fmt.Errorf("stock: %s has cost basis without a holding", symbol) + } + } + for symbol, qty := range p.Assets { + if qty <= 0 { + continue + } + if !isPositiveFiniteCost(p.CostBasis[symbol]) { + return fmt.Errorf("stock: holding %s has missing or invalid cost basis", symbol) + } + } + return nil +} + +func isPositiveFiniteCost(value float64) bool { + return value > 0 && !math.IsNaN(value) && !math.IsInf(value, 0) +} + // SavePortfolio persists the portfolio. func SavePortfolio(ctx context.Context, store Store, userID int64, p Portfolio) error { + if err := p.ValidateCostBasis(); err != nil { + return fmt.Errorf("stock: save portfolio %d: %w", userID, err) + } if err := store.Put(ctx, portfolioKey(userID), p); err != nil { return fmt.Errorf("stock: save portfolio %d: %w", userID, err) } diff --git a/internal/modules/stock/portfolio_test.go b/internal/modules/stock/portfolio_test.go index 2c014b0..28597b7 100644 --- a/internal/modules/stock/portfolio_test.go +++ b/internal/modules/stock/portfolio_test.go @@ -34,6 +34,7 @@ func TestSaveAndLoadRoundTrip(t *testing.T) { p, _ := LoadPortfolio(context.Background(), store, 42, 1) p.AddCurrency("VND", 5_000_000) p.AddAsset("TCB", 100) + p.CostBasis["TCB"] = 3_000_000 p.Meta.Invested = 5_000_000 if err := SavePortfolio(context.Background(), store, 42, p); err != nil { t.Fatalf("Save: %v", err) diff --git a/internal/modules/stock/startup.go b/internal/modules/stock/startup.go new file mode 100644 index 0000000..e2513e4 --- /dev/null +++ b/internal/modules/stock/startup.go @@ -0,0 +1,158 @@ +package stock + +import ( + "context" + "errors" + "fmt" + "sort" + "time" + + "github.com/tiennm99/miti99bot/internal/log" + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +const ( + CollectionName = "stock" + costBasisMigrationKey = "migration:stock-cost-basis-v1" + costBasisMigrationRetries = 5 +) + +// MigrationPriceFetcher supplies the current quotes used to seed legacy holdings. +type MigrationPriceFetcher interface { + FetchPrices(context.Context, []string) (map[string]float64, error) +} + +// InitStore ensures every existing stock holding has a cost basis. The scan runs +// every boot so the marker is an audit record, not permission to skip validation. +func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection, prices MigrationPriceFetcher) error { + docs := storage.Typed[Portfolio](portfolioColl) + system := systemstate.New(systemColl) + marker, markerExists, err := system.Get(ctx, costBasisMigrationKey) + if err != nil { + return fmt.Errorf("stock cost basis migration: read marker: %w", err) + } + keys, err := docs.List(ctx, "user:") + if err != nil { + return fmt.Errorf("stock cost basis migration: list portfolios: %w", err) + } + missing := map[string]bool{} + for _, key := range keys { + p, _, err := docs.Get(ctx, key) + if err != nil { + return fmt.Errorf("stock cost basis migration: read %s: %w", key, err) + } + if err := inspectLegacyPortfolio(p, missing); err != nil { + return fmt.Errorf("stock cost basis migration: %s: %w", key, err) + } + } + + symbols := sortedMissingSymbols(missing) + quotes := map[string]float64{} + if len(symbols) > 0 { + quotes, err = prices.FetchPrices(ctx, symbols) + if err != nil { + return fmt.Errorf("stock cost basis migration: fetch quotes: %w", err) + } + for _, symbol := range symbols { + if !isPositiveFiniteCost(quotes[symbol]) { + return fmt.Errorf("stock cost basis migration: no valid quote for %s", symbol) + } + } + } + + var migrated int64 + for index, key := range keys { + count, err := migrateLegacyPortfolio(ctx, docs, key, quotes) + if err != nil { + return err + } + migrated += int64(count) + if count > 0 { + log.Info("stock cost basis migrated", "portfolio", index+1, "total", len(keys), "positions", count) + } + } + now := time.Now().UnixMilli() + if markerExists && marker.Status == "completed" && migrated == 0 { + return nil + } + if !markerExists { + marker = systemstate.Record{Kind: "migration", Name: "stock cost basis v1", CompletedAt: now} + } + marker.Status = "completed" + marker.Count += migrated + marker.UpdatedAt = now + if err := system.Put(ctx, costBasisMigrationKey, marker); err != nil { + return fmt.Errorf("stock cost basis migration: write marker: %w", err) + } + return nil +} + +func inspectLegacyPortfolio(p Portfolio, missing map[string]bool) error { + for symbol, basis := range p.CostBasis { + if !isPositiveFiniteCost(basis) { + return fmt.Errorf("%s has invalid cost basis", symbol) + } + if p.Assets[symbol] <= 0 { + return fmt.Errorf("%s has cost basis without a holding", symbol) + } + } + for symbol, qty := range p.Assets { + if qty < 0 { + return fmt.Errorf("%s has negative quantity", symbol) + } + if qty == 0 { + continue + } + canonical, err := normalizeStockSymbol(symbol) + if err != nil || canonical != symbol { + return fmt.Errorf("%q is not a canonical ticker", symbol) + } + if _, ok := p.CostBasis[symbol]; !ok { + missing[symbol] = true + } + } + return nil +} + +func migrateLegacyPortfolio(ctx context.Context, docs storage.DocStore[Portfolio], key string, quotes map[string]float64) (int, error) { + for attempt := 0; attempt < costBasisMigrationRetries; attempt++ { + p, version, err := docs.Get(ctx, key) + if err != nil { + return 0, fmt.Errorf("stock cost basis migration: read %s: %w", key, err) + } + missing := map[string]bool{} + if err := inspectLegacyPortfolio(p, missing); err != nil { + return 0, fmt.Errorf("stock cost basis migration: %s: %w", key, err) + } + if len(missing) == 0 { + return 0, nil + } + if p.CostBasis == nil { + p.CostBasis = map[string]float64{} + } + for symbol := range missing { + quote := quotes[symbol] + basis := float64(p.Assets[symbol]) * quote + if !isPositiveFiniteCost(quote) || !isPositiveFiniteCost(basis) { + return 0, fmt.Errorf("stock cost basis migration: no cached valid quote for %s", symbol) + } + p.CostBasis[symbol] = basis + } + if err := docs.PutVersioned(ctx, key, version, p); err == nil { + return len(missing), nil + } else if !errors.Is(err, storage.ErrConflict) { + return 0, fmt.Errorf("stock cost basis migration: write %s: %w", key, err) + } + } + return 0, fmt.Errorf("stock cost basis migration: write %s: %w", key, storage.ErrConflict) +} + +func sortedMissingSymbols(missing map[string]bool) []string { + symbols := make([]string, 0, len(missing)) + for symbol := range missing { + symbols = append(symbols, symbol) + } + sort.Strings(symbols) + return symbols +} diff --git a/internal/modules/stock/startup_mongo_test.go b/internal/modules/stock/startup_mongo_test.go new file mode 100644 index 0000000..7361f7c --- /dev/null +++ b/internal/modules/stock/startup_mongo_test.go @@ -0,0 +1,60 @@ +package stock + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" + "github.com/tiennm99/miti99bot/internal/testutil/mongotest" +) + +var mongoTests mongotest.Manager + +func TestMain(m *testing.M) { + os.Exit(mongoTests.Run(m)) +} + +func TestInitStoreMigratesStockBasisInMongoDB(t *testing.T) { + ctx, portfolioColl, systemColl := setupMongoStockTest(t) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["TCB"] = 100 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + prices := &migrationStockPrices{quotes: map[string]float64{"TCB": 30_000}} + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore: %v", err) + } + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore second run: %v", err) + } + got, _, err := docs.Get(ctx, "user:7") + if err != nil || got.CostBasis["TCB"] != 3_000_000 { + t.Fatalf("portfolio=%+v err=%v", got, err) + } +} + +func setupMongoStockTest(t *testing.T) (context.Context, storage.Collection, storage.Collection) { + t.Helper() + uri := mongoTests.URI(t) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + client, err := storage.NewMongoClient(ctx, uri) + if err != nil { + t.Fatal(err) + } + db := client.Database(fmt.Sprintf("miti99bot_stock_test_%d", time.Now().UnixNano())) + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + _ = db.Drop(cleanupCtx) + _ = client.Disconnect(cleanupCtx) + }) + provider := storage.NewMongoProvider(db) + return ctx, provider.Collection(CollectionName), provider.Collection(systemstate.CollectionName) +} diff --git a/internal/modules/stock/startup_test.go b/internal/modules/stock/startup_test.go new file mode 100644 index 0000000..4d326c2 --- /dev/null +++ b/internal/modules/stock/startup_test.go @@ -0,0 +1,140 @@ +package stock + +import ( + "context" + "errors" + "testing" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +type migrationStockPrices struct { + quotes map[string]float64 + err error + calls int +} + +func (f *migrationStockPrices) FetchPrices(_ context.Context, _ []string) (map[string]float64, error) { + f.calls++ + return f.quotes, f.err +} + +func TestInitStoreMigratesLegacyStockBasisIdempotently(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + portfolioColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["TCB"] = 100 + legacy.Assets["FPT"] = 20 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + prices := &migrationStockPrices{quotes: map[string]float64{"TCB": 30_000, "FPT": 120_000}} + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore: %v", err) + } + got, _, err := docs.Get(ctx, "user:7") + if err != nil { + t.Fatal(err) + } + if got.CostBasis["TCB"] != 3_000_000 || got.CostBasis["FPT"] != 2_400_000 { + t.Fatalf("CostBasis = %#v", got.CostBasis) + } + if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil { + t.Fatalf("InitStore second run: %v", err) + } + if prices.calls != 1 { + t.Fatalf("quote calls = %d, want 1", prices.calls) + } + marker, _, err := storage.Typed[systemstate.Record](systemColl).Get(ctx, costBasisMigrationKey) + if err != nil || marker.Status != "completed" { + t.Fatalf("marker = %+v, err=%v", marker, err) + } +} + +func TestInitStoreStockRequiresCompleteQuotesBeforeWriting(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + portfolioColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + docs := storage.Typed[Portfolio](portfolioColl) + legacy := NewPortfolio(1) + legacy.Assets["TCB"] = 100 + if err := docs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + if err := InitStore(ctx, portfolioColl, systemColl, &migrationStockPrices{quotes: map[string]float64{}}); err == nil { + t.Fatal("InitStore succeeded without a complete quote set") + } + got, _, _ := docs.Get(ctx, "user:7") + if len(got.CostBasis) != 0 { + t.Fatalf("partial migration wrote basis: %#v", got.CostBasis) + } + if _, _, err := storage.Typed[systemstate.Record](systemColl).Get(ctx, costBasisMigrationKey); !errors.Is(err, storage.ErrNotFound) { + t.Fatalf("marker err = %v, want ErrNotFound", err) + } +} + +func TestStockWeightedAverageBasis(t *testing.T) { + p := NewPortfolio(1) + p.AddAsset("TCB", 100) + if err := p.AddCostBasis("TCB", 2_000_000); err != nil { + t.Fatal(err) + } + p.AddAsset("TCB", 50) + if err := p.AddCostBasis("TCB", 1_500_000); err != nil { + t.Fatal(err) + } + removed, err := p.RemoveCostBasis("TCB", 60, 150) + if err != nil { + t.Fatal(err) + } + if removed != 1_400_000 || p.CostBasis["TCB"] != 2_100_000 { + t.Fatalf("removed=%v remaining=%v", removed, p.CostBasis["TCB"]) + } +} + +type conflictOnceStockMigrationStore struct { + Store + conflicted bool +} + +func (s *conflictOnceStockMigrationStore) PutVersioned(ctx context.Context, key string, version int64, p Portfolio) error { + if !s.conflicted { + s.conflicted = true + return storage.ErrConflict + } + return s.Store.PutVersioned(ctx, key, version, p) +} + +func TestStockMigrationRetriesWriteConflictAndPreservesExistingBasis(t *testing.T) { + ctx := context.Background() + base := newStockStore() + legacy := NewPortfolio(1) + legacy.Assets["TCB"] = 100 + legacy.Assets["FPT"] = 10 + legacy.CostBasis["TCB"] = 2_500_000 + if err := base.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + store := &conflictOnceStockMigrationStore{Store: base} + count, err := migrateLegacyPortfolio(ctx, store, "user:7", map[string]float64{"FPT": 120_000}) + if err != nil || count != 1 || !store.conflicted { + t.Fatalf("count=%d conflicted=%v err=%v", count, store.conflicted, err) + } + got, _, _ := base.Get(ctx, "user:7") + if got.CostBasis["TCB"] != 2_500_000 || got.CostBasis["FPT"] != 1_200_000 { + t.Fatalf("CostBasis=%#v", got.CostBasis) + } +} + +func TestStockMigrationRejectsNoncanonicalLegacySymbol(t *testing.T) { + p := NewPortfolio(1) + p.Assets["tcb"] = 100 + if err := inspectLegacyPortfolio(p, map[string]bool{}); err == nil { + t.Fatal("inspectLegacyPortfolio accepted noncanonical symbol") + } +} diff --git a/internal/modules/stock/stats_test.go b/internal/modules/stock/stats_test.go index 92ae8a6..36d283e 100644 --- a/internal/modules/stock/stats_test.go +++ b/internal/modules/stock/stats_test.go @@ -43,6 +43,9 @@ func TestHandleStats_UsesSSIBatchPrices(t *testing.T) { p.AddAsset("MWG", 1800) p.AddAsset("TCB", 4200) p.AddAsset("FPT", 2300) + p.CostBasis["MWG"] = 108_000_000 + p.CostBasis["TCB"] = 105_000_000 + p.CostBasis["FPT"] = 230_000_000 if err := SavePortfolio(ctx, store, 7, p); err != nil { t.Fatalf("SavePortfolio: %v", err) } @@ -67,11 +70,12 @@ func TestHandleStats_UsesSSIBatchPrices(t *testing.T) { text := rb.LastSent().Text() for _, want := range []string{ - "MWG x1800 @ 70.000 VND = 126.000.000 VND", - "TCB x4200 @ 30.000 VND = 126.000.000 VND", - "FPT x2300 @ 120.000 VND = 276.000.000 VND", + "MWG x1800 | Avg 60.000 VND | Now 70.000 VND | Value 126.000.000 VND | Unrealized P&L +18.000.000 VND (+16.67%)", + "TCB x4200 | Avg 25.000 VND | Now 30.000 VND | Value 126.000.000 VND | Unrealized P&L +21.000.000 VND (+20.00%)", + "FPT x2300 | Avg 100.000 VND | Now 120.000 VND | Value 276.000.000 VND | Unrealized P&L +46.000.000 VND (+20.00%)", "Total value: 530.335.000 VND", - "P&L: -469.665.000 VND (-46.97%)", + "Unrealized P&L: +85.000.000 VND (+19.19%)", + "Account P&L: -469.665.000 VND (-46.97%)", } { if !strings.Contains(text, want) { t.Fatalf("stats missing %q in:\n%s", want, text) @@ -81,3 +85,17 @@ func TestHandleStats_UsesSSIBatchPrices(t *testing.T) { t.Fatalf("stats rendered missing prices:\n%s", text) } } + +func TestStockPortfolioReplyStaysWithinTelegramBudget(t *testing.T) { + positions := make([]string, 200) + for i := range positions { + positions[i] = strings.Repeat("position-data-", 20) + } + reply := boundedPortfolioReply([]string{"header"}, positions, []string{"summary"}) + if len(reply) > portfolioReplyLimit { + t.Fatalf("reply length = %d, limit = %d", len(reply), portfolioReplyLimit) + } + if !strings.Contains(reply, "position(s) omitted") || !strings.Contains(reply, "summary") { + t.Fatalf("bounded reply lost omission marker or summary: %q", reply) + } +}