tiennm99/miti99bot
1
0
Fork 0
mirror of https://github.com/tiennm99/miti99bot.git synced 2026-08-06 14:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(portfolio): nest asset positions

Browse Source
This commit is contained in:
tiennm99
2026-07-21 16:42:33 +07:00
parent 1c6e0ad8a0
commit 6fe91f0e4e
24 changed files with 830 additions and 979 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+15 -10
View File
@@ -53,25 +53,30 @@ notice and avoiding accidental repeated adjustments.
### Stock and coin P&L accounting
Stock and coin portfolios persist the total remaining cost basis for each open
position. Buys add their actual spend. Partial sells remove basis using the
Stock and coin portfolios embed each open position under `assets.<symbol>` with
`quantity`, total remaining `base`, and `dividendCheckedAt`. Stock cash is
stored directly as `vnd`; coin cash remains `usd`. Buys add their actual spend. Partial sells remove basis using the
weighted-average method and report realized P&L; full sells remove the position
and its basis. Stock share dividends add shares without adding cost, which
lowers the derived average price, while cash dividends do not change position
basis.
`/stock_portfolio` and `/coin_portfolio` show average entry price and
unrealized P&L for each priced position. `Account P&L` remains the broader
`/stock_portfolio` and `/coin_portfolio` show aligned monospace tables with
average entry price and unrealized P&L for each priced position. `Account P&L` remains the broader
account value minus all top-ups, so it also reflects realized proceeds,
dividend cash, and idle cash. If any current quote is unavailable, totals are
marked partial and numeric Account P&L is withheld.
On startup, enabled stock and coin modules scan every stored portfolio. Legacy
holdings without basis are initialized at that startup's current market quote,
giving them zero initial unrealized P&L. The migration uses optimistic writes,
records completion in the shared `system` collection, and still verifies the
invariant on every boot. Startup fails rather than accepting trades when a
required legacy quote or migration write is unavailable.
`dividendCheckedAt` is a future dividend-event cursor. It is initialized when a
position is first bought, preserved across later buys and sells, and advanced
when a stock dividend is recorded. A full exit removes the cursor; reopening
the position starts it again.
On startup, enabled stock and coin modules migrate the previous flat
`assets`/`costBasis` shape into embedded asset documents. Stock also migrates
`currency.VND` to `vnd`. The migration preserves the already-established basis,
uses optimistic writes, records completion in the shared `system` collection,
and verifies the new invariant every boot.
## Layout
cmd/server/main.go
+2 -2
View File
@@ -135,13 +135,13 @@ func main() {
}
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 {
if err := stock.InitStore(migrationCtx, provider.Collection(stock.CollectionName), provider.Collection(systemstate.CollectionName)); 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 {
if err := coin.InitStore(migrationCtx, provider.Collection(coin.CollectionName), provider.Collection(systemstate.CollectionName)); err != nil {
cancelMigrations()
log.Fatal("coin storage init failed", "err", err)
}
docs/deploy-coolify-selfhosted.md
+8 -6
View File
@@ -101,12 +101,14 @@ Successful GIF replies include the result behind Telegram spoiler formatting.
> counts and creates indexes on startup. Deleted legacy command rows are
> retained with `deleted: true`; `/stats` queries filter those rows from visible
> results. A historical `system` collection may remain in MongoDB with completed
> migration records. Stock and coin documents store `costBasis` by symbol. On
> every startup, enabled paper-trading modules verify that each positive holding
> has a valid basis; legacy holdings are initialized from current quotes before
> Telegram handlers are installed. The bot intentionally fails startup if a
> required quote or migration write fails. Completed records remain in `system`
> as audit history, but do not suppress later invariant scans.
> migration records. Stock stores cash as `vnd`; coin stores cash as `usd`.
> Both embed positions as
> `assets.<symbol>.{quantity,base,dividendCheckedAt}`. On every startup, enabled
> paper-trading modules verify this shape and migrate the previous flat
> `assets`/`costBasis` documents before Telegram handlers are installed. The bot
> intentionally fails startup if validation or a migration write fails.
> Completed records remain in `system` as audit history, but do not suppress
> later invariant scans.
## 2. Coolify
internal/modules/coin/handlers.go
+11 -15
View File
@@ -49,7 +49,8 @@ func (s *state) handleTopup(ctx context.Context, b *bot.Bot, update *models.Upda
return chathelper.Reply(ctx, b, update.Message, "Amount must be a positive finite USD number within the supported range.")
}
defer s.locks.Acquire(strconv.FormatInt(userID, 10))()
p, err := UpdatePortfolio(ctx, s.store, userID, s.now().UnixMilli(), func(p *Portfolio) error {
now := s.now().UnixMilli()
p, err := UpdatePortfolio(ctx, s.store, userID, now, func(p *Portfolio) error {
p.AddUSD(amount)
p.Meta.Invested += amount
return nil
@@ -94,17 +95,14 @@ func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update
}
defer s.locks.Acquire(strconv.FormatInt(userID, 10))()
var insufficientBalance *float64
p, err := UpdatePortfolio(ctx, s.store, userID, s.now().UnixMilli(), func(p *Portfolio) error {
now := s.now().UnixMilli()
p, err := UpdatePortfolio(ctx, s.store, userID, now, func(p *Portfolio) error {
ok, balance := p.DeductUSD(amount)
if !ok {
insufficientBalance = &balance
return errInsufficientUSD
}
if err := p.AddCostBasis(coin.Symbol, amount); err != nil {
return err
}
p.AddAsset(coin.Symbol, qty)
return nil
return p.BuyTicker(coin.Symbol, qty, amount, now)
})
if errors.Is(err, errInsufficientUSD) && insufficientBalance != nil {
return chathelper.Reply(ctx, b, update.Message,
@@ -154,19 +152,17 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat
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)
heldBefore := p.Assets[coin.Symbol].Quantity
_, removedBase, ok, sellErr := p.SellTicker(coin.Symbol, qty)
if !ok {
insufficientHeldQty = held
insufficientHeldQty = heldBefore
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
if sellErr != nil {
return sellErr
}
soldBasis = removedBase
p.AddUSD(amount)
return nil
})
internal/modules/coin/handlers_order_test.go
+1 -1
View File
@@ -20,7 +20,7 @@ func TestHandleBuyAcceptsCoinFirstOrder(t *testing.T) {
rb.AssertSentText(t, "Bought 0.0002 BTC")
p, _ := LoadPortfolio(ctx, s.store, 7, 999)
if p.USD != 990 || p.Assets["BTC"] != 0.0002 {
if p.USD != 990 || p.Assets["BTC"].Quantity != 0.0002 {
t.Fatalf("after coin-first buy = %+v", p)
}
}
internal/modules/coin/handlers_test.go
+9 -10
View File
@@ -153,11 +153,11 @@ func TestHandleBuyAndSell(t *testing.T) {
}
rb.AssertSentText(t, "Bought 0.01 BTC")
p, _ := LoadPortfolio(ctx, s.store, 7, 999)
if p.USD != 500 || p.Assets["BTC"] != 0.01 {
if p.USD != 500 || p.Assets["BTC"].Quantity != 0.01 {
t.Fatalf("after buy = %+v", p)
}
if p.CostBasis["BTC"] != 500 {
t.Fatalf("buy cost basis = %v, want 500", p.CostBasis["BTC"])
if p.Assets["BTC"].Base != 500 || p.Assets["BTC"].DividendCheckedAt != 123 {
t.Fatalf("buy asset = %+v", p.Assets["BTC"])
}
s.prices.(fakePriceFetcher).prices["BTC"] = CoinPrice{USD: 60_000, Source: "Binance"}
rb.Reset()
@@ -167,7 +167,7 @@ func TestHandleBuyAndSell(t *testing.T) {
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 != 1100 || len(p.Assets) != 0 || len(p.CostBasis) != 0 {
if p.USD != 1100 || len(p.Assets) != 0 {
t.Fatalf("after sell = %+v", p)
}
}
@@ -244,7 +244,7 @@ func TestHandleSellInsufficientCoinWithHoldings(t *testing.T) {
}
p, _ := LoadPortfolio(ctx, s.store, 7, 999)
if p.USD != 500 || p.Assets["BTC"] != 0.01 {
if p.USD != 500 || p.Assets["BTC"].Quantity != 0.01 {
t.Fatalf("portfolio mutated on failed sell = %+v", p)
}
}
@@ -334,7 +334,7 @@ func TestStatsWithAndWithoutPrice(t *testing.T) {
t.Fatalf("handleStats: %v", err)
}
text := rb.LastSent().Text()
for _, want := range []string{"Coin Account Summary", "BTC: 0.01", "(Binance)", "P&L:"} {
for _, want := range []string{"Coin Portfolio", "<pre>", "BTC", "0.01", "P&amp;L"} {
if !strings.Contains(text, want) {
t.Fatalf("stats missing %q in %q", want, text)
}
@@ -344,7 +344,7 @@ func TestStatsWithAndWithoutPrice(t *testing.T) {
if err := s.handleStats(ctx, rb.Bot, testutil.NewPrivateMessage(7, "/coin_portfolio")); err != nil {
t.Fatalf("handleStats no price: %v", err)
}
rb.AssertSentText(t, "price unavailable")
rb.AssertSentText(t, "N/A")
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())
}
@@ -354,8 +354,7 @@ 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
p.Assets["BTC"] = AssetPosition{Quantity: 2, Base: 1, DividendCheckedAt: 1}
if err := SavePortfolio(ctx, s.store, 7, p); err != nil {
t.Fatal(err)
}
@@ -364,7 +363,7 @@ func TestStatsTreatsOverflowedValuationAsUnavailable(t *testing.T) {
t.Fatal(err)
}
text := rb.LastSent().Text()
if !strings.Contains(text, "valuation unavailable") || !strings.Contains(text, "Account P&L: unavailable") {
if !strings.Contains(text, "N/A") || !strings.Contains(text, "Account P&amp;L") || !strings.Contains(text, "Unavailable") {
t.Fatalf("overflowed valuation was presented as complete: %q", text)
}
}
internal/modules/coin/portfolio.go
+96 -126
View File
@@ -1,26 +1,70 @@
package coin
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"math"
"strconv"
"go.mongodb.org/mongo-driver/v2/bson"
"github.com/tiennm99/miti99bot/internal/storage"
)
const coinDustEpsilon = 1e-9
const portfolioUpdateAttempts = 5
// Store is the coin module's typed portfolio store.
type Store = storage.DocStore[Portfolio]
type AssetPosition struct {
Quantity float64 `json:"quantity" bson:"quantity"`
Base float64 `json:"base" bson:"base"`
DividendCheckedAt int64 `json:"dividendCheckedAt" bson:"dividendCheckedAt"`
legacyQuantity bool
}
func (p *AssetPosition) UnmarshalJSON(data []byte) error {
data = bytes.TrimSpace(data)
if len(data) > 0 && data[0] == '{' {
type plain AssetPosition
return json.Unmarshal(data, (*plain)(p))
}
var quantity float64
if err := json.Unmarshal(data, &quantity); err != nil {
return fmt.Errorf("coin: decode legacy asset quantity: %w", err)
}
*p = AssetPosition{Quantity: quantity, legacyQuantity: true}
return nil
}
func (p *AssetPosition) UnmarshalBSONValue(valueType byte, data []byte) error {
raw := bson.RawValue{Type: bson.Type(valueType), Value: data}
if raw.Type == bson.TypeEmbeddedDocument {
type plain AssetPosition
return raw.Unmarshal((*plain)(p))
}
if quantity, ok := raw.DoubleOK(); ok {
*p = AssetPosition{Quantity: quantity, legacyQuantity: true}
return nil
}
if quantity, ok := raw.Int64OK(); ok {
*p = AssetPosition{Quantity: float64(quantity), legacyQuantity: true}
return nil
}
if quantity, ok := raw.Int32OK(); ok {
*p = AssetPosition{Quantity: float64(quantity), legacyQuantity: true}
return nil
}
return fmt.Errorf("coin: unsupported legacy asset BSON type %s", raw.Type)
}
type Portfolio struct {
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"`
USD float64 `json:"usd" bson:"usd"`
Assets map[string]AssetPosition `json:"assets" bson:"assets"`
Meta PortfolioMeta `json:"meta" bson:"meta"`
}
type PortfolioMeta struct {
@@ -29,16 +73,10 @@ type PortfolioMeta struct {
}
func NewPortfolio(now int64) Portfolio {
return Portfolio{
Assets: map[string]float64{},
CostBasis: map[string]float64{},
Meta: PortfolioMeta{CreatedAt: now},
}
return Portfolio{Assets: map[string]AssetPosition{}, Meta: PortfolioMeta{CreatedAt: now}}
}
func portfolioKey(userID int64) string {
return "user:" + strconv.FormatInt(userID, 10)
}
func portfolioKey(userID int64) string { return "user:" + strconv.FormatInt(userID, 10) }
func LoadPortfolio(ctx context.Context, store Store, userID int64, now int64) (Portfolio, error) {
p, _, err := loadPortfolioForUpdate(ctx, store, portfolioKey(userID), now)
@@ -50,7 +88,7 @@ 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 {
if err := p.Validate(); err != nil {
return fmt.Errorf("coin: save portfolio %d: %w", userID, err)
}
if err := store.Put(ctx, portfolioKey(userID), p); err != nil {
@@ -70,7 +108,7 @@ func UpdatePortfolio(ctx context.Context, store Store, userID int64, now int64,
return p, err
}
p.normalize()
if err := p.ValidateCostBasis(); err != nil {
if err := p.Validate(); err != nil {
return Portfolio{}, err
}
if err := store.PutVersioned(ctx, key, version, p); err == nil {
@@ -86,20 +124,13 @@ 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{}
p.Assets = map[string]AssetPosition{}
}
if p.Meta.CreatedAt == 0 {
p.Meta.CreatedAt = now
}
if err := p.ValidateCostBasis(); err != nil {
if err := p.Validate(); err != nil {
return Portfolio{}, 0, err
}
return p, version, nil
@@ -110,10 +141,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)
func (p Portfolio) Validate() error {
if math.IsNaN(p.USD) || math.IsInf(p.USD, 0) || p.USD < 0 {
return fmt.Errorf("coin: invalid USD balance")
}
for symbol, position := range p.Assets {
coin, err := ResolveCoinSymbol(symbol)
if err != nil || coin.Symbol != symbol || !isPositiveFinite(position.Quantity) ||
!isPositiveFinite(position.Base) || position.DividendCheckedAt <= 0 {
return fmt.Errorf("coin: %s has invalid position", symbol)
}
}
return nil
@@ -126,126 +162,60 @@ func (p *Portfolio) AddUSD(amount float64) {
func (p *Portfolio) DeductUSD(amount float64) (ok bool, balance float64) {
p.normalize()
balance = p.USD
if balance+coinDustEpsilon < amount {
return false, balance
if p.USD+coinDustEpsilon < amount {
return false, p.USD
}
p.USD = balance - amount
p.normalize()
p.USD = normalizeAmount(p.USD - amount)
return true, p.USD
}
func (p *Portfolio) AddAsset(symbol string, amount float64) {
func (p *Portfolio) BuyTicker(symbol string, quantity, base float64, now int64) error {
if !isPositiveFinite(quantity) || !isPositiveFinite(base) || now <= 0 {
return fmt.Errorf("coin: invalid purchase position")
}
if p.Assets == nil {
p.Assets = map[string]float64{}
p.Assets = map[string]AssetPosition{}
}
p.Assets[symbol] += amount
p.normalize()
}
func (p *Portfolio) AddCostBasis(symbol string, amount float64) error {
if !isPositiveFinite(amount) {
return fmt.Errorf("coin: invalid purchase cost basis")
position := p.Assets[symbol]
position.Quantity += quantity
position.Base += base
if !isPositiveFinite(position.Quantity) || !isPositiveFinite(position.Base) {
return fmt.Errorf("coin: position overflows")
}
if p.CostBasis == nil {
p.CostBasis = map[string]float64{}
if position.DividendCheckedAt == 0 {
position.DividendCheckedAt = now
}
next := p.CostBasis[symbol] + amount
if !isPositiveFinite(next) {
return fmt.Errorf("coin: cost basis overflows")
}
p.CostBasis[symbol] = next
p.Assets[symbol] = position
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")
func (p *Portfolio) SellTicker(symbol string, quantity float64) (remaining, soldBase float64, ok bool, err error) {
position, exists := p.Assets[symbol]
if !exists || !isPositiveFinite(quantity) || position.Quantity+coinDustEpsilon < quantity {
return position.Quantity, 0, false, nil
}
basis := p.CostBasis[symbol]
if !isPositiveFinite(basis) {
return 0, fmt.Errorf("coin: missing cost basis for %s", symbol)
remaining = normalizeAmount(position.Quantity - quantity)
if remaining == 0 {
delete(p.Assets, symbol)
return 0, position.Base, true, nil
}
if !holdingRemains {
delete(p.CostBasis, symbol)
return basis, nil
soldBase = position.Base * (quantity / position.Quantity)
position.Quantity = remaining
position.Base -= soldBase
if !isPositiveFinite(soldBase) || !isPositiveFinite(position.Base) {
return 0, 0, false, fmt.Errorf("coin: invalid remaining cost basis")
}
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{}
}
p.normalize()
held = p.Assets[symbol]
if held+coinDustEpsilon < amount {
return false, held
}
p.Assets[symbol] = held - amount
p.normalize()
return true, p.Assets[symbol]
p.Assets[symbol] = position
return remaining, soldBase, true, nil
}
func (p *Portfolio) normalize() {
p.USD = normalizeAmount(p.USD)
p.Meta.Invested = normalizeAmount(p.Meta.Invested)
if p.Assets == nil {
return
}
for symbol, amount := range p.Assets {
amount = normalizeAmount(amount)
if amount == 0 {
delete(p.Assets, symbol)
} else {
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 {
if math.IsNaN(n) || math.IsInf(n, 0) {
return 0
}
if math.Abs(n) < coinDustEpsilon {
if math.IsNaN(n) || math.IsInf(n, 0) || math.Abs(n) < coinDustEpsilon {
return 0
}
return n
internal/modules/coin/portfolio_test.go
+31 -72
View File
@@ -11,81 +11,46 @@ import (
func TestLoadPortfolioFirstTimeUser(t *testing.T) {
p, err := LoadPortfolio(context.Background(), newCoinStore(), 42, 123)
if err != nil {
t.Fatalf("LoadPortfolio: %v", err)
}
if p.USD != 0 || len(p.Assets) != 0 || p.Meta.CreatedAt != 123 {
t.Fatalf("portfolio = %+v", p)
if err != nil || p.USD != 0 || len(p.Assets) != 0 || p.Meta.CreatedAt != 123 {
t.Fatalf("portfolio=%+v err=%v", p, err)
}
}
func TestPortfolioBuySellMath(t *testing.T) {
func TestCoinBuySellMathAndCursor(t *testing.T) {
p := NewPortfolio(1)
p.AddUSD(1000)
p.Meta.Invested = 1000
if ok, bal := p.DeductUSD(250); !ok || bal != 750 {
t.Fatalf("DeductUSD ok=%v bal=%v", ok, bal)
if ok, balance := p.DeductUSD(250); !ok || balance != 750 {
t.Fatalf("balance=%v ok=%v", balance, ok)
}
p.AddAsset("BTC", 0.1)
if ok, held := p.DeductAsset("BTC", 0.04); !ok || math.Abs(held-0.06) > 1e-12 {
t.Fatalf("DeductAsset ok=%v held=%v", ok, held)
if err := p.BuyTicker("BTC", 0.1, 250, 10); err != nil {
t.Fatal(err)
}
if p.Assets["BTC"] <= 0 {
t.Fatalf("BTC holding missing: %+v", p.Assets)
if err := p.BuyTicker("BTC", 0.05, 150, 20); err != nil {
t.Fatal(err)
}
position := p.Assets["BTC"]
if position.DividendCheckedAt != 10 {
t.Fatalf("cursor=%d", position.DividendCheckedAt)
}
remaining, soldBase, ok, err := p.SellTicker("BTC", 0.06)
if err != nil || !ok || math.Abs(remaining-0.09) > 1e-12 || math.Abs(soldBase-160) > 1e-9 {
t.Fatalf("remaining=%v soldBase=%v ok=%v err=%v", remaining, soldBase, ok, err)
}
if p.Assets["BTC"].DividendCheckedAt != 10 {
t.Fatal("sell changed dividend cursor")
}
}
func TestDeductInsufficientBalances(t *testing.T) {
func TestCoinFullSellRemovesTicker(t *testing.T) {
p := NewPortfolio(1)
p.AddUSD(10)
p.AddAsset("ETH", 0.5)
if ok, bal := p.DeductUSD(11); ok || bal != 10 || p.USD != 10 {
t.Fatalf("DeductUSD ok=%v bal=%v p=%+v", ok, bal, p)
}
if ok, held := p.DeductAsset("ETH", 0.6); ok || held != 0.5 || p.Assets["ETH"] != 0.5 {
t.Fatalf("DeductAsset ok=%v held=%v p=%+v", ok, held, p)
_ = p.BuyTicker("BTC", 0.3, 12_000, 10)
_, soldBase, ok, err := p.SellTicker("BTC", 0.3)
if err != nil || !ok || math.Abs(soldBase-12_000) > 1e-9 || len(p.Assets) != 0 {
t.Fatalf("portfolio=%+v soldBase=%v ok=%v err=%v", p, soldBase, ok, err)
}
}
func TestDeductAssetDustCleanup(t *testing.T) {
p := NewPortfolio(1)
p.AddAsset("BTC", 0.1)
p.AddAsset("BTC", 0.2)
if ok, _ := p.DeductAsset("BTC", 0.3); !ok {
t.Fatal("DeductAsset ok=false")
}
if _, ok := p.Assets["BTC"]; ok {
t.Fatalf("dust key not removed: %+v", p.Assets)
}
}
func TestNormalizeAmountSpecialValues(t *testing.T) {
for _, n := range []float64{math.NaN(), math.Inf(1), math.Inf(-1), coinDustEpsilon / 2} {
if got := normalizeAmount(n); got != 0 {
t.Fatalf("normalizeAmount(%v) = %v, want 0", n, got)
}
}
}
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.
type conflictOnceStore struct {
Store
conflicted bool
@@ -105,30 +70,24 @@ func (s *conflictOnceStore) PutVersioned(ctx context.Context, key string, expect
}
func TestUpdatePortfolioRetriesAfterWriteConflict(t *testing.T) {
ctx := context.Background()
store := &conflictOnceStore{Store: newCoinStore()}
got, err := UpdatePortfolio(ctx, store, 7, 1, func(p *Portfolio) error {
got, err := UpdatePortfolio(context.Background(), store, 7, 1, func(p *Portfolio) error {
p.AddUSD(5)
return nil
})
if err != nil {
t.Fatalf("UpdatePortfolio: %v", err)
}
if got.USD != 15 {
t.Fatalf("USD = %v, want 15", got.USD)
if err != nil || got.USD != 15 {
t.Fatalf("USD=%v err=%v", got.USD, err)
}
}
func TestUpdatePortfolioMutateErrorDoesNotPersist(t *testing.T) {
ctx := context.Background()
store := newCoinStore()
_, err := UpdatePortfolio(ctx, store, 7, 1, func(p *Portfolio) error {
return errInsufficientUSD
})
_, err := UpdatePortfolio(ctx, store, 7, 1, func(*Portfolio) error { return errInsufficientUSD })
if !errors.Is(err, errInsufficientUSD) {
t.Fatalf("got %v, want errInsufficientUSD", err)
t.Fatalf("err=%v", err)
}
if _, _, err := store.Get(ctx, "user:7"); !errors.Is(err, storage.ErrNotFound) {
t.Fatalf("failed mutate must not persist, Get = %v", err)
t.Fatalf("Get err=%v", err)
}
}
internal/modules/coin/startup.go
+65 -106
View File
@@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"sort"
"time"
"github.com/tiennm99/miti99bot/internal/log"
@@ -13,64 +12,38 @@ import (
)
const (
CollectionName = "coin"
costBasisMigrationKey = "migration:coin-cost-basis-v1"
costBasisMigrationRetries = 5
CollectionName = "coin"
assetSchemaMarkerKey = "migration:coin-asset-schema-v2"
tickerMigrationRetries = 5
)
type MigrationPriceFetcher interface {
FetchUSD(context.Context, CoinSymbol) (CoinPrice, error)
type legacyPortfolio struct {
USD float64 `json:"usd" bson:"usd"`
Assets map[string]AssetPosition `json:"assets" bson:"assets"`
CostBasis map[string]float64 `json:"costBasis,omitempty" bson:"costBasis,omitempty"`
Meta PortfolioMeta `json:"meta" bson:"meta"`
}
// 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)
func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection) error {
docs := storage.Typed[legacyPortfolio](portfolioColl)
system := systemstate.New(systemColl)
marker, markerExists, err := system.Get(ctx, costBasisMigrationKey)
marker, markerExists, err := system.Get(ctx, assetSchemaMarkerKey)
if err != nil {
return fmt.Errorf("coin cost basis migration: read marker: %w", err)
return fmt.Errorf("coin asset schema 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)
return fmt.Errorf("coin asset schema 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)
changed, err := migrateAssetSchema(ctx, docs, key, time.Now().UnixMilli())
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)
if changed {
migrated++
log.Info("coin asset schema migrated", "portfolio", index+1, "total", len(keys))
}
}
now := time.Now().UnixMilli()
@@ -78,82 +51,68 @@ func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection
return nil
}
if !markerExists {
marker = systemstate.Record{Kind: "migration", Name: "coin cost basis v1", CompletedAt: now}
marker = systemstate.Record{Kind: "migration", Name: "coin asset schema v2", 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)
if err := system.Put(ctx, assetSchemaMarkerKey, marker); err != nil {
return fmt.Errorf("coin asset schema 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)
func migrateAssetSchema(ctx context.Context, docs storage.DocStore[legacyPortfolio], key string, now int64) (bool, error) {
for attempt := 0; attempt < tickerMigrationRetries; attempt++ {
doc, version, err := docs.Get(ctx, key)
if err != nil {
return false, fmt.Errorf("coin asset schema migration: read %s: %w", key, err)
}
if p.Assets[symbol] <= 0 {
return fmt.Errorf("%s has cost basis without a holding", symbol)
changed, err := doc.migrate(now)
if err != nil {
return false, fmt.Errorf("coin asset schema migration: %s: %w", key, err)
}
if !changed {
return false, nil
}
if err := docs.PutVersioned(ctx, key, version, doc); err == nil {
return true, nil
} else if !errors.Is(err, storage.ErrConflict) {
return false, fmt.Errorf("coin asset schema migration: write %s: %w", key, err)
}
}
for symbol, qty := range p.Assets {
if !isPositiveFinite(qty) {
if qty == 0 {
continue
}
return fmt.Errorf("%s has invalid quantity", symbol)
return false, fmt.Errorf("coin asset schema migration: write %s: %w", key, storage.ErrConflict)
}
func (p *legacyPortfolio) migrate(now int64) (bool, error) {
hasLegacyQuantity := false
for _, position := range p.Assets {
hasLegacyQuantity = hasLegacyQuantity || position.legacyQuantity
}
hasLegacy := p.CostBasis != nil || hasLegacyQuantity
if !hasLegacy {
return false, Portfolio{USD: p.USD, Assets: p.Assets, Meta: p.Meta}.Validate()
}
for symbol, position := range p.Assets {
if !position.legacyQuantity {
return false, fmt.Errorf("document mixes legacy and nested assets")
}
if position.Quantity == 0 {
delete(p.Assets, symbol)
continue
}
coin, err := ResolveCoinSymbol(symbol)
if err != nil || coin.Symbol != symbol {
return fmt.Errorf("%q is not a canonical coin symbol", symbol)
base := p.CostBasis[symbol]
if err != nil || coin.Symbol != symbol || !isPositiveFinite(position.Quantity) || !isPositiveFinite(base) {
return false, fmt.Errorf("invalid legacy position %q", symbol)
}
if _, ok := p.CostBasis[symbol]; !ok {
missing[symbol] = true
p.Assets[symbol] = AssetPosition{Quantity: position.Quantity, Base: base, DividendCheckedAt: now}
}
for symbol, base := range p.CostBasis {
if !isPositiveFinite(base) || p.Assets[symbol].Quantity <= 0 {
return false, fmt.Errorf("orphan or invalid legacy basis %q", symbol)
}
}
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
p.CostBasis = nil
return true, Portfolio{USD: p.USD, Assets: p.Assets, Meta: p.Meta}.Validate()
}
internal/modules/coin/startup_mongo_test.go
+20 -9
View File
@@ -7,6 +7,8 @@ import (
"testing"
"time"
"go.mongodb.org/mongo-driver/v2/bson"
"github.com/tiennm99/miti99bot/internal/storage"
"github.com/tiennm99/miti99bot/internal/systemstate"
"github.com/tiennm99/miti99bot/internal/testutil/mongotest"
@@ -20,23 +22,32 @@ func TestMain(m *testing.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 {
if err := storage.Typed[oldCoinPortfolio](portfolioColl).Put(ctx, "user:7", oldCoinPortfolio{
USD: 500, Assets: map[string]float64{"BTC": 0.25}, CostBasis: map[string]float64{"BTC": 25_000},
Meta: PortfolioMeta{CreatedAt: 1},
}); err != nil {
t.Fatal(err)
}
prices := &migrationCoinPrices{quotes: map[string]float64{"BTC": 100_000}}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("InitStore: %v", err)
}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("InitStore second run: %v", err)
}
got, _, err := docs.Get(ctx, "user:7")
if err != nil || got.CostBasis["BTC"] != 25_000 {
got, _, err := storage.Typed[Portfolio](portfolioColl).Get(ctx, "user:7")
if err != nil || got.USD != 500 || got.Assets["BTC"].Quantity != 0.25 || got.Assets["BTC"].Base != 25_000 {
t.Fatalf("portfolio=%+v err=%v", got, err)
}
rawColl, _ := storage.MongoCollection(portfolioColl)
var raw bson.M
if err := rawColl.FindOne(ctx, bson.M{"_id": "user:7"}).Decode(&raw); err != nil {
t.Fatal(err)
}
for _, legacyField := range []string{"costBasis", "tickers"} {
if _, exists := raw[legacyField]; exists {
t.Fatalf("legacy field %q remains in %#v", legacyField, raw)
}
}
}
func setupMongoCoinTest(t *testing.T) (context.Context, storage.Collection, storage.Collection) {
internal/modules/coin/startup_test.go
+36 -103
View File
@@ -2,140 +2,73 @@ 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
type oldCoinPortfolio struct {
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"`
}
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) {
func TestInitStoreMigratesCoinNestedAssets(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 {
if err := storage.Typed[oldCoinPortfolio](portfolioColl).Put(ctx, "user:7", oldCoinPortfolio{
USD: 500, Assets: map[string]float64{"BTC": 0.25}, CostBasis: map[string]float64{"BTC": 20_000},
Meta: PortfolioMeta{CreatedAt: 1},
}); 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 {
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("InitStore: %v", err)
}
got, _, err := docs.Get(ctx, "user:7")
got, err := LoadPortfolio(ctx, storage.Typed[Portfolio](portfolioColl), 7, 9)
if err != nil {
t.Fatal(err)
}
if got.CostBasis["BTC"] != 25_000 || got.CostBasis["ETH"] != 6_000 {
t.Fatalf("CostBasis = %#v", got.CostBasis)
position := got.Assets["BTC"]
if got.USD != 500 || position.Quantity != 0.25 || position.Base != 20_000 || position.DividendCheckedAt <= 0 {
t.Fatalf("portfolio=%+v", got)
}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
t.Fatalf("InitStore second run: %v", err)
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("second InitStore: %v", err)
}
if prices.calls != 2 {
t.Fatalf("quote calls = %d, want one per symbol on first run", prices.calls)
marker, exists, err := systemstate.New(systemColl).Get(ctx, assetSchemaMarkerKey)
if err != nil || !exists || marker.Count != 1 {
t.Fatalf("marker=%+v exists=%v err=%v", marker, exists, err)
}
}
func TestInitStoreCoinRequiresCompleteQuotesBeforeWriting(t *testing.T) {
func TestCoinMigrationRejectsMissingBasis(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 {
coll := provider.Collection(CollectionName)
if err := storage.Typed[oldCoinPortfolio](coll).Put(ctx, "user:7", oldCoinPortfolio{
Assets: map[string]float64{"BTC": 1}, CostBasis: map[string]float64{},
}); 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)
if err := InitStore(ctx, coll, provider.Collection(systemstate.CollectionName)); err == nil {
t.Fatal("InitStore accepted legacy holding without basis")
}
}
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)
func TestCoinSchemaMigrationRejectsMixedAssetShapes(t *testing.T) {
doc := legacyPortfolio{
Assets: map[string]AssetPosition{
"BTC": {Quantity: 1, Base: 10, DividendCheckedAt: 1},
"ETH": {Quantity: 1, legacyQuantity: true},
},
CostBasis: map[string]float64{"ETH": 2_000},
}
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)
if _, err := doc.migrate(123); err == nil {
t.Fatal("migration accepted mixed legacy and nested assets")
}
}
internal/modules/coin/views.go
+30 -32
View File
@@ -4,7 +4,6 @@ import (
"context"
"sort"
"strconv"
"strings"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
@@ -23,8 +22,7 @@ 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.")
}
header := []string{"Coin Account Summary", "USD: " + FormatUSD(p.USD)}
var positions []string
var positions [][]string
totalValue := p.USD
totalBasis := 0.0
missingPrice := false
@@ -37,56 +35,55 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda
fetchCtx, cancel := chathelper.FetchContext(ctx)
defer cancel()
for _, symbol := range sortedAssetSymbols(p.Assets) {
held := p.Assets[symbol]
basis := p.CostBasis[symbol]
held := p.Assets[symbol].Quantity
basis := p.Assets[symbol].Base
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 && isPositiveFinite(price.USD) {
value := held * price.USD
if !isPositiveFinite(value) || !isPositiveFinite(average) {
missingPrice = true
positions = append(positions, symbol+": "+FormatCoinQty(held)+" | valuation unavailable")
positions = append(positions, []string{symbol, FormatCoinQty(held), "N/A", "N/A", "N/A", "N/A"})
continue
}
totalValue += value
totalBasis += basis
line += " | Now " + FormatUSD(price.USD) + " (" + price.Source + ")" +
" | Value " + FormatUSD(value) + " | Unrealized P&L " + FormatPnLUSD(value, basis)
positions = append(positions, []string{symbol, FormatCoinQty(held), FormatUSD(average), FormatUSD(price.USD), FormatUSD(value), FormatPnLUSD(value, basis)})
} else {
log.Error("coin_fetch_price", "symbol", symbol, "err", err)
missingPrice = true
line += " | price unavailable"
positions = append(positions, []string{symbol, FormatCoinQty(held), FormatUSD(average), "N/A", "N/A", "N/A"})
}
} else {
missingPrice = true
line += " | price unavailable"
positions = append(positions, []string{symbol, FormatCoinQty(held), FormatUSD(average), "N/A", "N/A", "N/A"})
}
positions = append(positions, line)
}
var summary []string
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",
summary = [][]string{
{"USD", FormatUSD(p.USD)},
{"Priced value (partial)", FormatUSD(totalValue)},
{"Unrealized P&L (priced)", FormatPnLUSD(totalValue-p.USD, totalBasis)},
{"Invested", FormatUSD(p.Meta.Invested)},
{"Account P&L", "Unavailable"},
}
} 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),
summary = [][]string{
{"USD", FormatUSD(p.USD)},
{"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))
return chathelper.ReplyHTML(ctx, b, update.Message, portfolioTableReply("Coin Portfolio", positions, summary))
}
func sortedAssetSymbols(assets map[string]float64) []string {
func sortedAssetSymbols(assets map[string]AssetPosition) []string {
symbols := make([]string, 0, len(assets))
for symbol, amount := range assets {
if amount > 0 {
for symbol, position := range assets {
if position.Quantity > 0 {
symbols = append(symbols, symbol)
}
}
@@ -96,15 +93,16 @@ func sortedAssetSymbols(assets map[string]float64) []string {
const portfolioReplyLimit = 4000
func boundedPortfolioReply(header, positions, summary []string) string {
func portfolioTableReply(title string, positions, summary [][]string) string {
omitted := 0
for {
lines := append(append(append([]string{}, header...), positions...), summary...)
rows := append([][]string{}, positions...)
if omitted > 0 {
insertAt := len(header) + len(positions)
lines = append(lines[:insertAt], append([]string{"… " + strconv.Itoa(omitted) + " position(s) omitted"}, lines[insertAt:]...)...)
rows = append(rows, []string{"… " + strconv.Itoa(omitted) + " omitted"})
}
reply := strings.Join(lines, "\n")
reply := "<b>" + title + "</b>\n" +
chathelper.MonospaceTable([]string{"Ticker", "Qty", "Avg", "Now", "Value", "Unrealized P&L"}, rows) + "\n" +
chathelper.MonospaceTable([]string{"Metric", "Value"}, summary)
if len(reply) <= portfolioReplyLimit || len(positions) == 0 {
return reply
}
internal/modules/coin/views_reply_budget_test.go
+7 -3
View File
@@ -56,7 +56,7 @@ func TestHandleStatsDeliversReplyWhenUpstreamHangs(t *testing.T) {
t.Fatalf("handleStats took %v — fetch was not bounded below the reply reserve", elapsed)
}
sent := rb.LastSent().Text()
if !strings.Contains(sent, "Coin Account Summary") || !strings.Contains(sent, "price unavailable") {
if !strings.Contains(sent, "Coin Portfolio") || !strings.Contains(sent, "N/A") || !strings.Contains(sent, "<pre>") {
t.Fatalf("reply missing summary / degraded line; got:\n%s", sent)
}
}
@@ -66,11 +66,15 @@ func TestCoinPortfolioReplyStaysWithinTelegramBudget(t *testing.T) {
for i := range positions {
positions[i] = strings.Repeat("position-data-", 20)
}
reply := boundedPortfolioReply([]string{"header"}, positions, []string{"summary"})
rows := make([][]string, len(positions))
for index, position := range positions {
rows[index] = []string{position}
}
reply := portfolioTableReply("header", rows, [][]string{{"summary", "value"}})
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") {
if !strings.Contains(reply, "omitted") || !strings.Contains(reply, "summary") {
t.Fatalf("bounded reply lost omission marker or summary: %q", reply)
}
}
internal/modules/gold/portfolio_test.go
+3 -2
View File
@@ -234,8 +234,9 @@ 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 := stockPortfolio.BuyTicker("TCB", 100, 3_000_000, 1); err != nil {
t.Fatalf("buy stock: %v", err)
}
if err := stockmod.SavePortfolio(ctx, stockStore, 7, stockPortfolio); err != nil {
t.Fatalf("save stock: %v", err)
}
internal/modules/stock/handlers.go
+59 -56
View File
@@ -119,19 +119,20 @@ func (s *state) handleTopup(ctx context.Context, b *bot.Bot, update *models.Upda
defer s.locks.Acquire(strconv.FormatInt(userID, 10))()
p, err := LoadPortfolio(ctx, s.store, userID, s.now().UnixMilli())
now := s.now().UnixMilli()
p, err := LoadPortfolio(ctx, s.store, userID, now)
if err != nil {
log.Error("stock_load_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.")
}
p.AddCurrency("VND", amount)
p.AddVND(amount)
p.Meta.Invested += amount
if err := SavePortfolio(ctx, s.store, userID, p); err != nil {
log.Error("stock_save_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
}
return chathelper.Reply(ctx, b, update.Message,
"Topped up "+FormatVND(amount)+".\nBalance: "+FormatVND(p.Currency["VND"]))
"Topped up "+FormatVND(amount)+".\nBalance: "+FormatVND(p.VND))
}
func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update) error {
@@ -170,21 +171,21 @@ func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update
defer s.locks.Acquire(strconv.FormatInt(userID, 10))()
p, err := LoadPortfolio(ctx, s.store, userID, s.now().UnixMilli())
now := s.now().UnixMilli()
p, err := LoadPortfolio(ctx, s.store, userID, now)
if err != nil {
log.Error("stock_load_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.")
}
ok, balance := p.DeductCurrency("VND", cost)
ok, balance := p.DeductVND(cost)
if !ok {
return chathelper.Reply(ctx, b, update.Message,
"Insufficient VND. Need "+FormatVND(cost)+", have "+FormatVND(balance)+".")
}
if err := p.AddCostBasis(symbol, cost); err != nil {
if err := p.BuyTicker(symbol, qty, cost, now); 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)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
@@ -192,7 +193,7 @@ func (s *state) handleBuy(ctx context.Context, b *bot.Bot, update *models.Update
return chathelper.Reply(ctx, b, update.Message,
"Bought "+FormatStock(float64(qty))+" "+symbol+
" @ "+FormatVND(price)+"\nCost: "+FormatVND(cost)+
"\nRemaining: "+FormatVND(p.Currency["VND"]))
"\nRemaining: "+FormatVND(p.VND))
}
func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Update) error {
@@ -237,19 +238,17 @@ 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)
held, soldBasis, ok, basisErr := p.SellTicker(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)
if basisErr != nil {
log.Error("stock_remove_cost_basis", "user", userID, "ticker", symbol, "err", basisErr)
return chathelper.Reply(ctx, b, update.Message, "Portfolio cost basis is unavailable. Restart the bot or contact the owner.")
}
p.AddCurrency("VND", revenue)
p.AddVND(revenue)
if err := SavePortfolio(ctx, s.store, userID, p); err != nil {
log.Error("stock_save_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
@@ -258,7 +257,7 @@ func (s *state) handleSell(ctx context.Context, b *bot.Bot, update *models.Updat
"Sold "+FormatStock(float64(qty))+" "+symbol+
" @ "+FormatVND(price)+"\nRevenue: "+FormatVND(revenue)+
"\nRealized P&L: "+FormatPnL(revenue, soldBasis)+
"\nRemaining: "+FormatVND(p.Currency["VND"]))
"\nRemaining: "+FormatVND(p.VND))
}
func (s *state) handleCashDividend(ctx context.Context, b *bot.Bot, update *models.Update) error {
@@ -293,7 +292,7 @@ func (s *state) handleCashDividend(ctx context.Context, b *bot.Bot, update *mode
log.Error("stock_load_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.")
}
held := p.Assets[symbol]
held := p.Assets[symbol].Quantity
if held <= 0 {
return chathelper.Reply(ctx, b, update.Message,
"You don't hold any "+symbol+" to receive a cash dividend.")
@@ -302,11 +301,13 @@ func (s *state) handleCashDividend(ctx context.Context, b *bot.Bot, update *mode
if err != nil {
return chathelper.Reply(ctx, b, update.Message, "Dividend amount is too large.")
}
balance, err := checkedVNDBalance(p.Currency["VND"], total)
balance, err := checkedVNDBalance(p.VND, total)
if err != nil {
return chathelper.Reply(ctx, b, update.Message, "Dividend amount is too large.")
}
p.Currency["VND"] = balance
if err := p.ApplyDividend(symbol, held, balance, s.now().UnixMilli()); err != nil {
return chathelper.Reply(ctx, b, update.Message, "Could not update dividend checkpoint. Try again later.")
}
if err := SavePortfolio(ctx, s.store, userID, p); err != nil {
log.Error("stock_save_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
@@ -349,7 +350,7 @@ func (s *state) handleShareDividend(ctx context.Context, b *bot.Bot, update *mod
log.Error("stock_load_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.")
}
held := p.Assets[symbol]
held := p.Assets[symbol].Quantity
if held <= 0 {
return chathelper.Reply(ctx, b, update.Message,
"You don't hold any "+symbol+" to receive a share dividend.")
@@ -367,7 +368,9 @@ func (s *state) handleShareDividend(ctx context.Context, b *bot.Bot, update *mod
if err != nil {
return chathelper.Reply(ctx, b, update.Message, "Share dividend is too large.")
}
p.Assets[symbol] = finalHolding
if err := p.ApplyDividend(symbol, finalHolding, p.VND, s.now().UnixMilli()); err != nil {
return chathelper.Reply(ctx, b, update.Message, "Could not update dividend checkpoint. Try again later.")
}
if err := SavePortfolio(ctx, s.store, userID, p); err != nil {
log.Error("stock_save_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
@@ -412,7 +415,7 @@ func (s *state) handleDividend(ctx context.Context, b *bot.Bot, update *models.U
log.Error("stock_load_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not load portfolio. Try again later.")
}
held := p.Assets[symbol]
held := p.Assets[symbol].Quantity
if held <= 0 {
return chathelper.Reply(ctx, b, update.Message,
"You don't hold any "+symbol+" to receive a dividend.")
@@ -429,13 +432,14 @@ func (s *state) handleDividend(ctx context.Context, b *bot.Bot, update *models.U
if err != nil {
return chathelper.Reply(ctx, b, update.Message, "Share dividend is too large.")
}
balance, err := checkedVNDBalance(p.Currency["VND"], total)
balance, err := checkedVNDBalance(p.VND, total)
if err != nil {
return chathelper.Reply(ctx, b, update.Message, "Dividend amount is too large.")
}
p.Currency["VND"] = balance
p.Assets[symbol] = finalHolding
if err := p.ApplyDividend(symbol, finalHolding, balance, s.now().UnixMilli()); err != nil {
return chathelper.Reply(ctx, b, update.Message, "Could not update dividend checkpoint. Try again later.")
}
if err := SavePortfolio(ctx, s.store, userID, p); err != nil {
log.Error("stock_save_portfolio", "user", userID, "err", err)
return chathelper.Reply(ctx, b, update.Message, "Could not save portfolio. Try again later.")
@@ -462,10 +466,8 @@ 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.")
}
header := []string{"📊 Portfolio Summary", "VND: " + FormatVND(p.Currency["VND"])}
var positions []string
totalValue := 0.0
totalValue += p.Currency["VND"]
var positions [][]string
totalValue := p.VND
totalBasis := 0.0
missingPrice := false
@@ -475,15 +477,14 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda
qty int64
}
var heldList []held
for sym, qty := range p.Assets {
if qty > 0 {
heldList = append(heldList, held{sym, qty})
for symbol, position := range p.Assets {
if position.Quantity > 0 {
heldList = append(heldList, held{symbol, position.Quantity})
}
}
sort.Slice(heldList, func(i, j int) bool { return heldList[i].symbol < heldList[j].symbol })
if len(heldList) > 0 {
header = append(header, "", "Stocks:")
fetchCtx, cancel := chathelper.FetchContext(ctx)
defer cancel()
@@ -498,55 +499,57 @@ func (s *state) handleStats(ctx context.Context, b *bot.Bot, update *models.Upda
for _, h := range heldList {
price := prices[h.symbol]
basis := p.CostBasis[h.symbol]
basis := p.Assets[h.symbol].Base
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")
positions = append(positions, []string{h.symbol, FormatStock(float64(h.qty)), FormatVND(average), "N/A", "N/A", "N/A"})
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")
positions = append(positions, []string{h.symbol, FormatStock(float64(h.qty)), "N/A", "N/A", "N/A", "N/A"})
continue
}
totalValue += 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))
positions = append(positions, []string{h.symbol, FormatStock(float64(h.qty)), FormatVND(average), FormatVND(price), FormatVND(val), FormatPnL(val, basis)})
}
}
var summary []string
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")
summary = [][]string{
{"VND", FormatVND(p.VND)},
{"Priced value (partial)", FormatVND(totalValue)},
{"Unrealized P&L (priced)", FormatPnL(totalValue-p.VND, totalBasis)},
{"Invested", FormatVND(p.Meta.Invested)},
{"Account P&L", "Unavailable"},
}
} 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))
summary = [][]string{
{"VND", FormatVND(p.VND)},
{"Total value", FormatVND(totalValue)},
{"Invested", FormatVND(p.Meta.Invested)},
{"Unrealized P&L", FormatPnL(totalValue-p.VND, totalBasis)},
{"Account P&L", FormatPnL(totalValue, p.Meta.Invested)},
}
}
return chathelper.Reply(ctx, b, update.Message, boundedPortfolioReply(header, positions, summary))
return chathelper.ReplyHTML(ctx, b, update.Message, portfolioTableReply("Stock Portfolio", positions, summary))
}
const portfolioReplyLimit = 4000
func boundedPortfolioReply(header, positions, summary []string) string {
func portfolioTableReply(title string, positions, summary [][]string) string {
omitted := 0
for {
lines := append(append(append([]string{}, header...), positions...), summary...)
rows := append([][]string{}, positions...)
if omitted > 0 {
insertAt := len(header) + len(positions)
lines = append(lines[:insertAt], append([]string{" … " + strconv.Itoa(omitted) + " position(s) omitted"}, lines[insertAt:]...)...)
rows = append(rows, []string{"… " + strconv.Itoa(omitted) + " omitted"})
}
reply := strings.Join(lines, "\n")
reply := "<b>" + title + "</b>\n" +
chathelper.MonospaceTable([]string{"Ticker", "Qty", "Avg", "Now", "Value", "Unrealized P&L"}, rows) + "\n" +
chathelper.MonospaceTable([]string{"Metric", "Value"}, summary)
if len(reply) <= portfolioReplyLimit || len(positions) == 0 {
return reply
}
internal/modules/stock/handlers_test.go
+18 -19
View File
@@ -97,7 +97,7 @@ func TestHandleBuyAndPartialSellTracksCostBasisAndRealizedPnL(t *testing.T) {
t.Fatal(err)
}
p, err := LoadPortfolio(ctx, s.store, 7, 999)
if err != nil || p.CostBasis["TCB"] != 3_000_000 {
if err != nil || p.Assets["TCB"].Base != 3_000_000 {
t.Fatalf("after buy=%+v err=%v", p, err)
}
price = 36_000
@@ -107,7 +107,7 @@ func TestHandleBuyAndPartialSellTracksCostBasisAndRealizedPnL(t *testing.T) {
}
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 {
if err != nil || p.Assets["TCB"].Quantity != 60 || p.Assets["TCB"].Base != 1_800_000 {
t.Fatalf("after sell=%+v err=%v", p, err)
}
}
@@ -188,7 +188,7 @@ func TestMutableHandlersRejectExtraArgs(t *testing.T) {
if err != nil {
t.Fatalf("LoadPortfolio: %v", err)
}
if p.Currency["VND"] != 0 || len(p.Assets) != 0 {
if p.VND != 0 || len(p.Assets) != 0 {
t.Fatalf("invalid commands mutated portfolio: %+v", p)
}
}
@@ -243,9 +243,8 @@ func (s *countingPortfolioStore) Put(ctx context.Context, id string, p Portfolio
func seedStockPortfolio(t *testing.T, store Store, userID int64, held int64, balance float64) {
t.Helper()
p := NewPortfolio(123)
p.Assets["TCB"] = held
p.CostBasis["TCB"] = float64(held) * 30_000
p.Currency["VND"] = balance
p.Assets["TCB"] = AssetPosition{Quantity: held, Base: float64(held) * 30_000, DividendCheckedAt: 100}
p.VND = balance
if err := SavePortfolio(context.Background(), store, userID, p); err != nil {
t.Fatalf("seed portfolio: %v", err)
}
@@ -267,11 +266,11 @@ func TestHandleCashDividendAllowsRepeatedManualAdjustments(t *testing.T) {
if err != nil {
t.Fatalf("load portfolio: %v", err)
}
if got, want := p.Currency["VND"], float64(418000); got != want {
if got, want := p.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"])
if p.Assets["TCB"].Base != 139*30_000 || p.Assets["TCB"].DividendCheckedAt != 123 {
t.Fatalf("cash dividend position: %+v", p.Assets["TCB"])
}
}
@@ -290,7 +289,7 @@ func TestHandleCashDividendRejectsInexactBalanceSum(t *testing.T) {
t.Fatalf("store writes = %d, want 0", store.puts)
}
p, _ := LoadPortfolio(ctx, base, 7, 999)
if p.Assets["TCB"] != 1 || p.Currency["VND"] != 1 {
if p.Assets["TCB"].Quantity != 1 || p.VND != 1 {
t.Fatalf("portfolio changed: %+v", p)
}
rb.AssertSentText(t, "Dividend amount is too large.")
@@ -307,11 +306,11 @@ func TestHandleShareDividendPreservesRatioAndFloors(t *testing.T) {
t.Fatalf("handleShareDividend: %v", err)
}
p, _ := LoadPortfolio(ctx, store, 7, 999)
if got, want := p.Assets["TCB"], int64(152); got != want {
if got, want := p.Assets["TCB"].Quantity, 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"])
if p.Assets["TCB"].Base != 139*30_000 || p.Assets["TCB"].DividendCheckedAt != 123 {
t.Fatalf("share dividend position: %+v", p.Assets["TCB"])
}
rb.AssertSentText(t, "Share dividend (100:10): +13 TCB")
}
@@ -345,8 +344,8 @@ func TestHandleShareDividendRejectsZeroEntitlement(t *testing.T) {
t.Fatalf("store writes = %d, want 0", store.puts)
}
p, _ := LoadPortfolio(ctx, base, 7, 999)
if p.Assets["TCB"] != 9 {
t.Fatalf("holding changed to %d", p.Assets["TCB"])
if p.Assets["TCB"].Quantity != 9 {
t.Fatalf("holding changed to %d", p.Assets["TCB"].Quantity)
}
rb.AssertSentText(t, "Minimum holding: 10")
}
@@ -379,7 +378,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 || p.CostBasis["TCB"] != 139*30_000 {
if p.Assets["TCB"].Quantity != 152 || p.VND != 209500 || p.Assets["TCB"].Base != 139*30_000 || p.Assets["TCB"].DividendCheckedAt != 123 {
t.Fatalf("portfolio = %+v", p)
}
rb.AssertSentText(t, "Dividend for TCB (100:10)")
@@ -401,7 +400,7 @@ func TestHandleCombinedDividendRejectsInexactBalanceSum(t *testing.T) {
t.Fatalf("store writes = %d, want 0", store.puts)
}
p, _ := LoadPortfolio(ctx, base, 7, 999)
if p.Assets["TCB"] != 1 || p.Currency["VND"] != 1 {
if p.Assets["TCB"].Quantity != 1 || p.VND != 1 {
t.Fatalf("portfolio changed: %+v", p)
}
rb.AssertSentText(t, "Dividend amount is too large.")
@@ -418,7 +417,7 @@ func TestHandleCombinedDividendCreditsCashWhenSharesRoundToZero(t *testing.T) {
t.Fatalf("handleDividend: %v", err)
}
p, _ := LoadPortfolio(ctx, store, 7, 999)
if p.Assets["TCB"] != 9 || p.Currency["VND"] != 13600 {
if p.Assets["TCB"].Quantity != 9 || p.VND != 13600 {
t.Fatalf("portfolio = %+v", p)
}
rb.AssertSentText(t, "Shares: +0")
@@ -436,7 +435,7 @@ func TestDividendSaveFailureLeavesStoredPortfolioUnchanged(t *testing.T) {
t.Fatalf("handleDividend: %v", err)
}
p, _ := LoadPortfolio(ctx, base, 7, 999)
if p.Assets["TCB"] != 139 || p.Currency["VND"] != 1000 {
if p.Assets["TCB"].Quantity != 139 || p.VND != 1000 {
t.Fatalf("stored portfolio changed: %+v", p)
}
rb.AssertSentText(t, "Could not save portfolio")
internal/modules/stock/portfolio.go
+133 -128
View File
@@ -1,69 +1,92 @@
package stock
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"math"
"strconv"
"go.mongodb.org/mongo-driver/v2/bson"
"github.com/tiennm99/miti99bot/internal/storage"
)
// Store is the stock module's typed portfolio store.
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.
// AssetPosition keeps the complete persisted state for one stock ticker.
// Base is total remaining VND cost, not average price. DividendCheckedAt is
// the cursor for future dividend-event discovery.
type AssetPosition struct {
Quantity int64 `json:"quantity" bson:"quantity"`
Base float64 `json:"base" bson:"base"`
DividendCheckedAt int64 `json:"dividendCheckedAt" bson:"dividendCheckedAt"`
legacyQuantity bool
}
func (p *AssetPosition) UnmarshalJSON(data []byte) error {
data = bytes.TrimSpace(data)
if len(data) > 0 && data[0] == '{' {
type plain AssetPosition
return json.Unmarshal(data, (*plain)(p))
}
var quantity int64
if err := json.Unmarshal(data, &quantity); err != nil {
return fmt.Errorf("stock: decode legacy asset quantity: %w", err)
}
*p = AssetPosition{Quantity: quantity, legacyQuantity: true}
return nil
}
func (p *AssetPosition) UnmarshalBSONValue(valueType byte, data []byte) error {
raw := bson.RawValue{Type: bson.Type(valueType), Value: data}
if raw.Type == bson.TypeEmbeddedDocument {
type plain AssetPosition
return raw.Unmarshal((*plain)(p))
}
if quantity, ok := raw.Int64OK(); ok {
*p = AssetPosition{Quantity: quantity, legacyQuantity: true}
return nil
}
if quantity, ok := raw.Int32OK(); ok {
*p = AssetPosition{Quantity: int64(quantity), legacyQuantity: true}
return nil
}
return fmt.Errorf("stock: unsupported legacy asset BSON type %s", raw.Type)
}
type Portfolio struct {
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"`
VND float64 `json:"vnd" bson:"vnd"`
Assets map[string]AssetPosition `json:"assets" bson:"assets"`
Meta PortfolioMeta `json:"meta" bson:"meta"`
}
// PortfolioMeta tracks invested cost basis for P&L. CreatedAt is purely
// informational (ms-epoch when the portfolio first existed).
type PortfolioMeta struct {
Invested float64 `json:"invested" bson:"invested"`
CreatedAt int64 `json:"createdAt" bson:"createdAt"`
}
// NewPortfolio returns an empty starting state. Currency map seeded with VND=0
// so deductCurrency on a fresh user reports "0 balance" cleanly instead of
// nil-map panics.
func NewPortfolio(now int64) Portfolio {
return Portfolio{
Currency: map[string]float64{"VND": 0},
Assets: map[string]int64{},
CostBasis: map[string]float64{},
Meta: PortfolioMeta{Invested: 0, CreatedAt: now},
}
return Portfolio{Assets: map[string]AssetPosition{}, Meta: PortfolioMeta{CreatedAt: now}}
}
func portfolioKey(userID int64) string {
return "user:" + strconv.FormatInt(userID, 10)
}
// LoadPortfolio reads from store; returns an empty portfolio on first-time use.
// Defensively initialises nil maps so callers never need a nil check.
func LoadPortfolio(ctx context.Context, store Store, userID int64, now int64) (Portfolio, error) {
p, _, err := store.Get(ctx, portfolioKey(userID))
switch {
case err == nil:
// Repair any nils from older / partial saves — defence in depth.
if p.Currency == nil {
p.Currency = map[string]float64{"VND": 0}
} else if _, ok := p.Currency["VND"]; !ok {
p.Currency["VND"] = 0
}
if p.Assets == nil {
p.Assets = map[string]int64{}
p.Assets = map[string]AssetPosition{}
}
if p.CostBasis == nil {
p.CostBasis = map[string]float64{}
if p.Meta.CreatedAt == 0 {
p.Meta.CreatedAt = now
}
if err := p.ValidateCostBasis(); err != nil {
if err := p.Validate(); err != nil {
return Portfolio{}, err
}
return p, nil
@@ -74,69 +97,8 @@ 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 {
if err := p.Validate(); err != nil {
return fmt.Errorf("stock: save portfolio %d: %w", userID, err)
}
if err := store.Put(ctx, portfolioKey(userID), p); err != nil {
@@ -145,52 +107,95 @@ func SavePortfolio(ctx context.Context, store Store, userID int64, p Portfolio)
return nil
}
// AddCurrency credits the currency balance.
func (p *Portfolio) AddCurrency(currency string, amount float64) {
if p.Currency == nil {
p.Currency = map[string]float64{}
func (p Portfolio) Validate() error {
if math.IsNaN(p.VND) || math.IsInf(p.VND, 0) || p.VND < 0 {
return fmt.Errorf("stock: invalid VND balance")
}
p.Currency[currency] += amount
for symbol, position := range p.Assets {
canonical, err := normalizeStockSymbol(symbol)
if err != nil || canonical != symbol {
return fmt.Errorf("stock: invalid ticker %q", symbol)
}
if position.Quantity <= 0 || !isPositiveFiniteCost(position.Base) || position.DividendCheckedAt <= 0 {
return fmt.Errorf("stock: %s has invalid position", symbol)
}
}
return nil
}
// DeductCurrency debits the currency balance. Returns false + the current
// balance when insufficient — caller renders the user-facing error.
func (p *Portfolio) DeductCurrency(currency string, amount float64) (ok bool, balance float64) {
if p.Currency == nil {
p.Currency = map[string]float64{}
}
balance = p.Currency[currency]
if balance < amount {
return false, balance
}
p.Currency[currency] = balance - amount
return true, p.Currency[currency]
func (p *Portfolio) AddVND(amount float64) {
p.VND += amount
}
// AddAsset credits the share holding.
func (p *Portfolio) AddAsset(symbol string, qty int64) {
func (p *Portfolio) DeductVND(amount float64) (ok bool, balance float64) {
if p.VND < amount {
return false, p.VND
}
p.VND -= amount
return true, p.VND
}
// BuyTicker adds quantity and basis. DividendCheckedAt is set only when opening
// a new position so later buys cannot skip events after the original buy.
func (p *Portfolio) BuyTicker(symbol string, quantity int64, base float64, now int64) error {
if quantity <= 0 || !isPositiveFiniteCost(base) || now <= 0 {
return fmt.Errorf("stock: invalid purchase position")
}
if p.Assets == nil {
p.Assets = map[string]int64{}
p.Assets = map[string]AssetPosition{}
}
p.Assets[symbol] += qty
position := p.Assets[symbol]
if position.Quantity > math.MaxInt64-quantity {
return fmt.Errorf("stock: quantity overflows")
}
position.Quantity += quantity
position.Base += base
if !isPositiveFiniteCost(position.Base) {
return fmt.Errorf("stock: cost basis overflows")
}
if position.DividendCheckedAt == 0 {
position.DividendCheckedAt = now
}
p.Assets[symbol] = position
return nil
}
// DeductAsset debits the share holding. Returns false + held when caller asks
// for more than they own. Removes the key when balance hits zero so the
// portfolio doesn't accumulate empty entries.
func (p *Portfolio) DeductAsset(symbol string, qty int64) (ok bool, held int64) {
if p.Assets == nil {
p.Assets = map[string]int64{}
// SellTicker removes proportional weighted-average basis while preserving the
// dividend cursor. A full exit removes the entire ticker document.
func (p *Portfolio) SellTicker(symbol string, quantity int64) (remaining int64, soldBase float64, ok bool, err error) {
position, exists := p.Assets[symbol]
if !exists || position.Quantity < quantity || quantity <= 0 {
return position.Quantity, 0, false, nil
}
held = p.Assets[symbol]
if held < qty {
return false, held
}
remaining := held - qty
if remaining == 0 {
if quantity == position.Quantity {
delete(p.Assets, symbol)
} else {
p.Assets[symbol] = remaining
return 0, position.Base, true, nil
}
return true, remaining
soldBase = position.Base * (float64(quantity) / float64(position.Quantity))
position.Quantity -= quantity
position.Base -= soldBase
if !isPositiveFiniteCost(soldBase) || !isPositiveFiniteCost(position.Base) {
return 0, 0, false, fmt.Errorf("stock: invalid remaining cost basis")
}
p.Assets[symbol] = position
return position.Quantity, soldBase, true, nil
}
func (p *Portfolio) ApplyDividend(symbol string, quantity int64, vnd float64, now int64) error {
position, ok := p.Assets[symbol]
if !ok || position.Quantity <= 0 {
return fmt.Errorf("stock: ticker position not found")
}
if quantity < position.Quantity || !isPositiveFiniteCost(position.Base) || now <= 0 {
return fmt.Errorf("stock: invalid dividend position")
}
position.Quantity = quantity
position.DividendCheckedAt = now
p.Assets[symbol] = position
p.VND = vnd
return nil
}
func isPositiveFiniteCost(value float64) bool {
return value > 0 && !math.IsNaN(value) && !math.IsInf(value, 0)
}
internal/modules/stock/portfolio_test.go
+44 -63
View File
@@ -7,91 +7,72 @@ import (
"github.com/tiennm99/miti99bot/internal/storage"
)
// newStockStore returns a fresh in-memory typed portfolio store for tests.
func newStockStore() Store {
return storage.Typed[Portfolio](storage.NewMemoryProvider().Collection("stock"))
return storage.Typed[Portfolio](storage.NewMemoryProvider().Collection(CollectionName))
}
func TestLoadPortfolio_FirstTimeUser(t *testing.T) {
store := newStockStore()
p, err := LoadPortfolio(context.Background(), store, 42, 1234567890)
func TestLoadPortfolioFirstTimeUser(t *testing.T) {
p, err := LoadPortfolio(context.Background(), newStockStore(), 42, 1234567890)
if err != nil {
t.Fatalf("LoadPortfolio: %v", err)
t.Fatal(err)
}
if p.Currency["VND"] != 0 {
t.Errorf("VND seeded: got %v, want 0", p.Currency["VND"])
}
if p.Assets == nil {
t.Error("Assets is nil")
}
if p.Meta.CreatedAt != 1234567890 {
t.Errorf("CreatedAt: got %d, want 1234567890", p.Meta.CreatedAt)
if p.VND != 0 || p.Assets == nil || p.Meta.CreatedAt != 1234567890 {
t.Fatalf("portfolio=%+v", p)
}
}
func TestSaveAndLoadRoundTrip(t *testing.T) {
ctx := context.Background()
store := newStockStore()
p, _ := LoadPortfolio(context.Background(), store, 42, 1)
p.AddCurrency("VND", 5_000_000)
p.AddAsset("TCB", 100)
p.CostBasis["TCB"] = 3_000_000
p := NewPortfolio(1)
p.AddVND(5_000_000)
if err := p.BuyTicker("TCB", 100, 3_000_000, 10); err != nil {
t.Fatal(err)
}
p.Meta.Invested = 5_000_000
if err := SavePortfolio(context.Background(), store, 42, p); err != nil {
t.Fatalf("Save: %v", err)
if err := SavePortfolio(ctx, store, 42, p); err != nil {
t.Fatal(err)
}
got, err := LoadPortfolio(context.Background(), store, 42, 999) // CreatedAt should NOT be reset
got, err := LoadPortfolio(ctx, store, 42, 999)
if err != nil {
t.Fatalf("Load: %v", err)
t.Fatal(err)
}
if got.Currency["VND"] != 5_000_000 {
t.Errorf("VND: got %v, want 5000000", got.Currency["VND"])
}
if got.Assets["TCB"] != 100 {
t.Errorf("TCB: got %d, want 100", got.Assets["TCB"])
}
if got.Meta.Invested != 5_000_000 {
t.Errorf("Invested: got %v, want 5000000", got.Meta.Invested)
}
if got.Meta.CreatedAt != 1 {
t.Errorf("CreatedAt: got %d, want 1 (load must NOT overwrite existing)", got.Meta.CreatedAt)
position := got.Assets["TCB"]
if got.VND != 5_000_000 || position.Quantity != 100 || position.Base != 3_000_000 || position.DividendCheckedAt != 10 || got.Meta.CreatedAt != 1 {
t.Fatalf("portfolio=%+v", got)
}
}
func TestAddDeductCurrency(t *testing.T) {
p := NewPortfolio(0)
p.AddCurrency("VND", 1000)
p.AddCurrency("VND", 500)
if p.Currency["VND"] != 1500 {
t.Errorf("after add: got %v, want 1500", p.Currency["VND"])
func TestBuyPreservesDividendCheckedAtAndSellUsesWeightedBasis(t *testing.T) {
p := NewPortfolio(1)
if err := p.BuyTicker("TCB", 100, 2_000_000, 10); err != nil {
t.Fatal(err)
}
ok, bal := p.DeductCurrency("VND", 600)
if !ok || bal != 900 {
t.Errorf("deduct 600: ok=%v bal=%v, want ok=true bal=900", ok, bal)
if err := p.BuyTicker("TCB", 50, 1_500_000, 20); err != nil {
t.Fatal(err)
}
ok, bal = p.DeductCurrency("VND", 9999)
if ok || bal != 900 {
t.Errorf("deduct over balance: ok=%v bal=%v, want ok=false bal=900 (unchanged)", ok, bal)
position := p.Assets["TCB"]
if position.DividendCheckedAt != 10 {
t.Fatalf("additional buy changed dividend cursor to %d", position.DividendCheckedAt)
}
remaining, soldBase, ok, err := p.SellTicker("TCB", 60)
if err != nil || !ok || remaining != 90 || soldBase != 1_400_000 {
t.Fatalf("remaining=%d soldBase=%v ok=%v err=%v", remaining, soldBase, ok, err)
}
position = p.Assets["TCB"]
if position.Base != 2_100_000 || position.DividendCheckedAt != 10 {
t.Fatalf("position=%+v", position)
}
}
func TestAddDeductAsset(t *testing.T) {
p := NewPortfolio(0)
p.AddAsset("TCB", 10)
p.AddAsset("TCB", 5)
if p.Assets["TCB"] != 15 {
t.Errorf("TCB after add: got %d, want 15", p.Assets["TCB"])
func TestDividendAdvancesCursorWithoutChangingBase(t *testing.T) {
p := NewPortfolio(1)
_ = p.BuyTicker("TCB", 100, 3_000_000, 10)
if err := p.ApplyDividend("TCB", 110, 500_000, 30); err != nil {
t.Fatal(err)
}
ok, held := p.DeductAsset("TCB", 3)
if !ok || held != 12 {
t.Errorf("deduct 3: ok=%v held=%v, want ok=true held=12", ok, held)
}
ok, _ = p.DeductAsset("TCB", 999)
if ok {
t.Error("deduct over holdings: should fail")
}
// Final deduction removes key entirely.
p.DeductAsset("TCB", 12)
if _, present := p.Assets["TCB"]; present {
t.Error("zero-balance asset should be removed from map")
position := p.Assets["TCB"]
if position.Quantity != 110 || position.Base != 3_000_000 || position.DividendCheckedAt != 30 || p.VND != 500_000 {
t.Fatalf("portfolio=%+v", p)
}
}
internal/modules/stock/startup.go
+105 -100
View File
@@ -4,7 +4,7 @@ import (
"context"
"errors"
"fmt"
"sort"
"math"
"time"
"github.com/tiennm99/miti99bot/internal/log"
@@ -13,63 +13,42 @@ import (
)
const (
CollectionName = "stock"
costBasisMigrationKey = "migration:stock-cost-basis-v1"
costBasisMigrationRetries = 5
CollectionName = "stock"
assetSchemaMarkerKey = "migration:stock-asset-schema-v2"
tickerMigrationRetries = 5
)
// MigrationPriceFetcher supplies the current quotes used to seed legacy holdings.
type MigrationPriceFetcher interface {
FetchPrices(context.Context, []string) (map[string]float64, error)
type legacyPortfolio struct {
VND float64 `json:"vnd,omitempty" bson:"vnd,omitempty"`
Currency map[string]float64 `json:"currency,omitempty" bson:"currency,omitempty"`
Assets map[string]AssetPosition `json:"assets" bson:"assets"`
CostBasis map[string]float64 `json:"costBasis,omitempty" bson:"costBasis,omitempty"`
Meta PortfolioMeta `json:"meta" bson:"meta"`
}
// 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)
// InitStore migrates the old currency/assets/costBasis shape into the nested
// asset schema before handlers can load portfolios. It remains an every-boot
// invariant scan; the system marker is audit history, not a bypass.
func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection) error {
docs := storage.Typed[legacyPortfolio](portfolioColl)
system := systemstate.New(systemColl)
marker, markerExists, err := system.Get(ctx, costBasisMigrationKey)
marker, markerExists, err := system.Get(ctx, assetSchemaMarkerKey)
if err != nil {
return fmt.Errorf("stock cost basis migration: read marker: %w", err)
return fmt.Errorf("stock asset schema 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)
return fmt.Errorf("stock asset schema 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)
changed, err := migrateAssetSchema(ctx, docs, key, time.Now().UnixMilli())
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)
if changed {
migrated++
log.Info("stock asset schema migrated", "portfolio", index+1, "total", len(keys))
}
}
now := time.Now().UnixMilli()
@@ -77,82 +56,108 @@ func InitStore(ctx context.Context, portfolioColl, systemColl storage.Collection
return nil
}
if !markerExists {
marker = systemstate.Record{Kind: "migration", Name: "stock cost basis v1", CompletedAt: now}
marker = systemstate.Record{Kind: "migration", Name: "stock asset schema v2", 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)
if err := system.Put(ctx, assetSchemaMarkerKey, marker); err != nil {
return fmt.Errorf("stock asset schema 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)
func migrateAssetSchema(ctx context.Context, docs storage.DocStore[legacyPortfolio], key string, now int64) (bool, error) {
for attempt := 0; attempt < tickerMigrationRetries; attempt++ {
doc, version, err := docs.Get(ctx, key)
if err != nil {
return false, fmt.Errorf("stock asset schema migration: read %s: %w", key, err)
}
if p.Assets[symbol] <= 0 {
return fmt.Errorf("%s has cost basis without a holding", symbol)
changed, err := doc.migrate(now)
if err != nil {
return false, fmt.Errorf("stock asset schema migration: %s: %w", key, err)
}
if !changed {
return false, nil
}
if err := docs.PutVersioned(ctx, key, version, doc); err == nil {
return true, nil
} else if !errors.Is(err, storage.ErrConflict) {
return false, fmt.Errorf("stock asset schema migration: write %s: %w", key, err)
}
}
for symbol, qty := range p.Assets {
if qty < 0 {
return fmt.Errorf("%s has negative quantity", symbol)
return false, fmt.Errorf("stock asset schema migration: write %s: %w", key, storage.ErrConflict)
}
func (p *legacyPortfolio) migrate(now int64) (bool, error) {
if !p.hasLegacyFields() {
return false, p.currentPortfolio().Validate()
}
if err := p.migrateVND(); err != nil {
return false, err
}
if err := p.migrateAssets(now); err != nil {
return false, err
}
p.Currency = nil
p.CostBasis = nil
return true, p.currentPortfolio().Validate()
}
func (p *legacyPortfolio) hasLegacyFields() bool {
if p.Currency != nil || p.CostBasis != nil {
return true
}
for _, position := range p.Assets {
if position.legacyQuantity {
return true
}
if qty == 0 {
}
return false
}
func (p *legacyPortfolio) migrateVND() error {
if math.IsNaN(p.VND) || math.IsInf(p.VND, 0) || p.VND < 0 {
return fmt.Errorf("invalid VND balance")
}
for currency, balance := range p.Currency {
if currency != "VND" && balance != 0 {
return fmt.Errorf("unsupported legacy currency %q", currency)
}
}
if legacyVND := p.Currency["VND"]; legacyVND != 0 {
if p.VND != 0 && p.VND != legacyVND {
return fmt.Errorf("conflicting VND balances")
}
p.VND = legacyVND
}
return nil
}
func (p *legacyPortfolio) migrateAssets(now int64) error {
for symbol, position := range p.Assets {
if !position.legacyQuantity {
return fmt.Errorf("document mixes legacy and nested assets")
}
if position.Quantity == 0 {
delete(p.Assets, symbol)
continue
}
canonical, err := normalizeStockSymbol(symbol)
if err != nil || canonical != symbol {
return fmt.Errorf("%q is not a canonical ticker", symbol)
base := p.CostBasis[symbol]
if err != nil || canonical != symbol || position.Quantity < 0 || !isPositiveFiniteCost(base) {
return fmt.Errorf("invalid legacy position %q", symbol)
}
if _, ok := p.CostBasis[symbol]; !ok {
missing[symbol] = true
p.Assets[symbol] = AssetPosition{Quantity: position.Quantity, Base: base, DividendCheckedAt: now}
}
for symbol, base := range p.CostBasis {
if !isPositiveFiniteCost(base) || p.Assets[symbol].Quantity <= 0 {
return fmt.Errorf("orphan or invalid legacy basis %q", symbol)
}
}
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
func (p *legacyPortfolio) currentPortfolio() Portfolio {
return Portfolio{VND: p.VND, Assets: p.Assets, Meta: p.Meta}
}
internal/modules/stock/startup_mongo_test.go
+20 -9
View File
@@ -7,6 +7,8 @@ import (
"testing"
"time"
"go.mongodb.org/mongo-driver/v2/bson"
"github.com/tiennm99/miti99bot/internal/storage"
"github.com/tiennm99/miti99bot/internal/systemstate"
"github.com/tiennm99/miti99bot/internal/testutil/mongotest"
@@ -20,23 +22,32 @@ func TestMain(m *testing.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 {
if err := storage.Typed[oldStockPortfolio](portfolioColl).Put(ctx, "user:7", oldStockPortfolio{
Currency: map[string]float64{"VND": 500_000}, Assets: map[string]int64{"TCB": 100},
CostBasis: map[string]float64{"TCB": 3_000_000}, Meta: PortfolioMeta{CreatedAt: 1},
}); err != nil {
t.Fatal(err)
}
prices := &migrationStockPrices{quotes: map[string]float64{"TCB": 30_000}}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("InitStore: %v", err)
}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
if err := InitStore(ctx, portfolioColl, systemColl); 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 {
got, _, err := storage.Typed[Portfolio](portfolioColl).Get(ctx, "user:7")
if err != nil || got.VND != 500_000 || got.Assets["TCB"].Quantity != 100 || got.Assets["TCB"].Base != 3_000_000 {
t.Fatalf("portfolio=%+v err=%v", got, err)
}
rawColl, _ := storage.MongoCollection(portfolioColl)
var raw bson.M
if err := rawColl.FindOne(ctx, bson.M{"_id": "user:7"}).Decode(&raw); err != nil {
t.Fatal(err)
}
for _, legacyField := range []string{"currency", "costBasis", "tickers"} {
if _, exists := raw[legacyField]; exists {
t.Fatalf("legacy field %q remains in %#v", legacyField, raw)
}
}
}
func setupMongoStockTest(t *testing.T) (context.Context, storage.Collection, storage.Collection) {
internal/modules/stock/startup_test.go
+46 -91
View File
@@ -2,139 +2,94 @@ 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
type oldStockPortfolio struct {
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"`
}
func (f *migrationStockPrices) FetchPrices(_ context.Context, _ []string) (map[string]float64, error) {
f.calls++
return f.quotes, f.err
}
func TestInitStoreMigratesLegacyStockBasisIdempotently(t *testing.T) {
func TestInitStoreMigratesStockNestedAssetsAndVND(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 {
oldDocs := storage.Typed[oldStockPortfolio](portfolioColl)
if err := oldDocs.Put(ctx, "user:7", oldStockPortfolio{
Currency: map[string]float64{"VND": 2_000_000},
Assets: map[string]int64{"TCB": 100}, CostBasis: map[string]float64{"TCB": 3_000_000},
Meta: PortfolioMeta{CreatedAt: 1},
}); 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 {
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("InitStore: %v", err)
}
got, _, err := docs.Get(ctx, "user:7")
got, err := LoadPortfolio(ctx, storage.Typed[Portfolio](portfolioColl), 7, 9)
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)
position := got.Assets["TCB"]
if got.VND != 2_000_000 || position.Quantity != 100 || position.Base != 3_000_000 || position.DividendCheckedAt <= 0 {
t.Fatalf("portfolio=%+v", got)
}
if err := InitStore(ctx, portfolioColl, systemColl, prices); err != nil {
t.Fatalf("InitStore second run: %v", err)
if err := InitStore(ctx, portfolioColl, systemColl); err != nil {
t.Fatalf("second InitStore: %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)
marker, exists, err := systemstate.New(systemColl).Get(ctx, assetSchemaMarkerKey)
if err != nil || !exists || marker.Status != "completed" || marker.Count != 1 {
t.Fatalf("marker=%+v exists=%v err=%v", marker, exists, err)
}
}
func TestInitStoreStockRequiresCompleteQuotesBeforeWriting(t *testing.T) {
func TestStockMigrationRejectsMissingBasis(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 {
coll := provider.Collection(CollectionName)
if err := storage.Typed[oldStockPortfolio](coll).Put(ctx, "user:7", oldStockPortfolio{
Currency: map[string]float64{"VND": 1}, Assets: map[string]int64{"TCB": 100}, CostBasis: map[string]float64{},
}); 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)
if err := InitStore(ctx, coll, provider.Collection(systemstate.CollectionName)); err == nil {
t.Fatal("InitStore accepted legacy holding without basis")
}
}
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
type conflictOnceSchemaStore struct {
storage.DocStore[legacyPortfolio]
conflicted bool
}
func (s *conflictOnceStockMigrationStore) PutVersioned(ctx context.Context, key string, version int64, p Portfolio) error {
func (s *conflictOnceSchemaStore) PutVersioned(ctx context.Context, key string, version int64, value legacyPortfolio) error {
if !s.conflicted {
s.conflicted = true
return storage.ErrConflict
}
return s.Store.PutVersioned(ctx, key, version, p)
return s.DocStore.PutVersioned(ctx, key, version, value)
}
func TestStockMigrationRetriesWriteConflictAndPreservesExistingBasis(t *testing.T) {
func TestStockSchemaMigrationRetriesConflict(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 {
provider := storage.NewMemoryProvider()
coll := provider.Collection(CollectionName)
if err := storage.Typed[oldStockPortfolio](coll).Put(ctx, "user:7", oldStockPortfolio{
Currency: map[string]float64{"VND": 1},
Assets: map[string]int64{"TCB": 10},
CostBasis: map[string]float64{"TCB": 300_000},
}); 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")
docs := storage.Typed[legacyPortfolio](coll)
store := &conflictOnceSchemaStore{DocStore: docs}
changed, err := migrateAssetSchema(ctx, store, "user:7", 123)
if err != nil || !changed || !store.conflicted {
t.Fatalf("changed=%v conflicted=%v err=%v", changed, store.conflicted, err)
}
}
internal/modules/stock/stats_test.go
+21 -16
View File
@@ -38,14 +38,11 @@ func TestHandleStats_UsesSSIBatchPrices(t *testing.T) {
store := newStockStore()
p := NewPortfolio(now.UnixMilli())
p.Currency["VND"] = 2335000
p.VND = 2335000
p.Meta.Invested = 1000000000
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
_ = p.BuyTicker("MWG", 1800, 108_000_000, 1)
_ = p.BuyTicker("TCB", 4200, 105_000_000, 1)
_ = p.BuyTicker("FPT", 2300, 230_000_000, 1)
if err := SavePortfolio(ctx, store, 7, p); err != nil {
t.Fatalf("SavePortfolio: %v", err)
}
@@ -70,18 +67,22 @@ func TestHandleStats_UsesSSIBatchPrices(t *testing.T) {
text := rb.LastSent().Text()
for _, want := range []string{
"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",
"Unrealized P&L: +85.000.000 VND (+19.19%)",
"Account P&L: -469.665.000 VND (-46.97%)",
"<pre>",
"Ticker",
"MWG",
"126.000.000 VND",
"Total value",
"530.335.000 VND",
"Unrealized P&amp;L",
"+85.000.000 VND (+19.19%)",
"Account P&amp;L",
"-469.665.000 VND (-46.97%)",
} {
if !strings.Contains(text, want) {
t.Fatalf("stats missing %q in:\n%s", want, text)
}
}
if strings.Contains(text, "(no price)") {
if strings.Contains(text, "N/A") {
t.Fatalf("stats rendered missing prices:\n%s", text)
}
}
@@ -91,11 +92,15 @@ func TestStockPortfolioReplyStaysWithinTelegramBudget(t *testing.T) {
for i := range positions {
positions[i] = strings.Repeat("position-data-", 20)
}
reply := boundedPortfolioReply([]string{"header"}, positions, []string{"summary"})
rows := make([][]string, len(positions))
for index, position := range positions {
rows[index] = []string{position}
}
reply := portfolioTableReply("header", rows, [][]string{{"summary", "value"}})
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") {
if !strings.Contains(reply, "omitted") || !strings.Contains(reply, "summary") {
t.Fatalf("bounded reply lost omission marker or summary: %q", reply)
}
}
internal/modules/util/chathelper/chathelper.go
+40
View File
@@ -6,6 +6,7 @@ package chathelper
import (
"context"
"html"
"math"
"strconv"
"strings"
@@ -15,6 +16,45 @@ import (
"github.com/go-telegram/bot/models"
)
// MonospaceTable renders an HTML-safe, left-aligned table for Telegram's
// <pre> mode. Callers should send the result with ReplyHTML.
func MonospaceTable(headers []string, rows [][]string) string {
widths := make([]int, len(headers))
for index, header := range headers {
widths[index] = len([]rune(header))
}
for _, row := range rows {
for index := range headers {
if index < len(row) && len([]rune(row[index])) > widths[index] {
widths[index] = len([]rune(row[index]))
}
}
}
var lines []string
appendRow := func(row []string) {
cells := make([]string, len(headers))
for index := range headers {
if index < len(row) {
cells[index] = row[index]
}
if index < len(headers)-1 {
cells[index] += strings.Repeat(" ", widths[index]-len([]rune(cells[index])))
}
}
lines = append(lines, strings.Join(cells, " "))
}
appendRow(headers)
separator := make([]string, len(headers))
for index := range headers {
separator[index] = strings.Repeat("-", widths[index])
}
appendRow(separator)
for _, row := range rows {
appendRow(row)
}
return "<pre>" + html.EscapeString(strings.Join(lines, "\n")) + "</pre>"
}
// SubjectFor returns the identity key per-module state should be scoped by:
// group/supergroup → chat ID (shared game state), otherwise → user ID.
// Returns "" when no usable id is present (caller should reply with a
internal/modules/util/chathelper/chathelper_test.go
+10
View File
@@ -2,6 +2,7 @@ package chathelper
import (
"context"
"strings"
"testing"
"time"
@@ -268,3 +269,12 @@ func TestWinRate(t *testing.T) {
}
}
}
func TestMonospaceTableAlignsAndEscapesHTML(t *testing.T) {
got := MonospaceTable([]string{"Asset", "Qty"}, [][]string{{"<TCB>", "10"}, {"BTC", "2"}})
for _, want := range []string{"<pre>", "Asset", "&lt;TCB&gt;", "BTC 2", "</pre>"} {
if !strings.Contains(got, want) {
t.Fatalf("table missing %q in %q", want, got)
}
}
}