stock-analysis/internal/backtest/account_test.go

package backtest

import (
	"math"
	"testing"
)

// mockPrice 返回一个简单的价格函数，所有标的固定价格。
func mockPrice(price float64) priceFn {
	return func(code, date string) (float64, bool) {
		return price, true
	}
}

// mockPriceMap 返回按 code 返回不同价格的价格函数。
func mockPriceMap(prices map[string]float64) priceFn {
	return func(code, date string) (float64, bool) {
		p, ok := prices[code]
		return p, ok
	}
}

func TestNewAccount(t *testing.T) {
	params := DefaultParams()
	params.InitialCash = 200000
	a := NewAccount(params)
	if a.cash != 200000 {
		t.Errorf("cash: expected 200000, got %f", a.cash)
	}
	if a.equity != 200000 {
		t.Errorf("equity: expected 200000, got %f", a.equity)
	}
	if len(a.holdings) != 0 {
		t.Errorf("holdings: expected empty, got %d", len(a.holdings))
	}
}

func TestMarkToMarket_NoHoldings(t *testing.T) {
	a := NewAccount(DefaultParams())
	a.MarkToMarket("2020-01-02", mockPrice(100))
	if a.equity != 100000 {
		t.Errorf("equity: expected 100000, got %f", a.equity)
	}
	if len(a.curve) != 1 {
		t.Errorf("curve: expected 1 point, got %d", len(a.curve))
	}
	if a.curve[0].Return != 0 {
		t.Errorf("return: expected 0, got %f", a.curve[0].Return)
	}
}

func TestMarkToMarket_WithHoldings(t *testing.T) {
	a := NewAccount(DefaultParams())
	// 手动设置持仓：100 张，每张 100 元
	a.holdings["113001"] = 100
	a.MarkToMarket("2020-01-02", mockPrice(100))
	// equity = cash(100000) + 100*100 = 110000
	expected := 100000.0 + 100*100
	if math.Abs(a.equity-expected) > 0.001 {
		t.Errorf("equity: expected %f, got %f", expected, a.equity)
	}
	// return = 110000/100000 - 1 = 0.1
	if math.Abs(a.curve[0].Return-0.1) > 0.001 {
		t.Errorf("return: expected 0.1, got %f", a.curve[0].Return)
	}
}

func TestRebalance_BuyOnly(t *testing.T) {
	params := Params{
		InitialCash: 100000,
		Commission:  0,
		Slippage:    0,
	}
	a := NewAccount(params)

	targets := []Target{
		{Code: "113001", Weight: 0.5},
		{Code: "113002", Weight: 0.5},
	}

	prices := map[string]float64{"113001": 100, "113002": 120}
	a.Rebalance("2020-01-02", targets, mockPriceMap(prices))

	// 总权益 100000
	// 113001: 50000 / (100*10) = 50 手 = 500 张, cost = 500*100 = 50000
	// 113002: 50000 / (120*10) = 41.67 → 41 手 = 410 张, cost = 410*120 = 49200
	// 剩余现金 = 100000 - 50000 - 49200 = 800
	if a.holdings["113001"] != 500 {
		t.Errorf("113001 holdings: expected 500, got %f", a.holdings["113001"])
	}
	if a.holdings["113002"] != 410 {
		t.Errorf("113002 holdings: expected 410, got %f", a.holdings["113002"])
	}
	if math.Abs(a.cash-800) > 0.01 {
		t.Errorf("cash: expected 800, got %f", a.cash)
	}
}

func TestRebalance_NoTargets(t *testing.T) {
	a := NewAccount(DefaultParams())
	a.holdings["113001"] = 100
	a.Rebalance("2020-01-02", nil, mockPrice(100))
	// 无目标 → 全部卖出
	if a.holdings["113001"] != 0 {
		t.Errorf("holdings should be sold, got %f", a.holdings["113001"])
	}
}

func TestRebalance_SellOnly(t *testing.T) {
	params := Params{
		InitialCash: 100000,
		Commission:  0,
		Slippage:    0,
	}
	a := NewAccount(params)
	a.holdings["113001"] = 500
	// 目标：只持有 113002，权重 1.0
	targets := []Target{{Code: "113002", Weight: 1.0}}
	prices := map[string]float64{"113001": 100, "113002": 100}
	a.Rebalance("2020-01-02", targets, mockPriceMap(prices))

	// 113001 应被全部卖出
	if _, ok := a.holdings["113001"]; ok {
		t.Errorf("113001 should be sold")
	}
	// 113002 应被买入
	if a.holdings["113002"] <= 0 {
		t.Errorf("113002 should be bought")
	}
}

func TestLots(t *testing.T) {
	if lots(5.7) != 5 {
		t.Errorf("lots(5.7): expected 5, got %f", lots(5.7))
	}
	if lots(10.0) != 10 {
		t.Errorf("lots(10.0): expected 10, got %f", lots(10.0))
	}
	if lots(0.3) != 0 {
		t.Errorf("lots(0.3): expected 0, got %f", lots(0.3))
	}
}

func TestEquity_InitialValue(t *testing.T) {
	a := NewAccount(DefaultParams())
	if a.Equity() != 100000 {
		t.Errorf("initial equity: expected 100000, got %f", a.Equity())
	}
}

func TestCurve_Accumulation(t *testing.T) {
	a := NewAccount(DefaultParams())
	a.MarkToMarket("2020-01-02", mockPrice(100))
	a.MarkToMarket("2020-01-03", mockPrice(100))
	c := a.Curve()
	if len(c) != 2 {
		t.Errorf("curve length: expected 2, got %d", len(c))
	}
}