"""Tests for the capital-allocation simulation engine.""" from __future__ import annotations import unittest from app.simulation import Candidate, allocate_capital class AllocationTest(unittest.TestCase): def make_candidates(self): # dividend + high score return [ Candidate("A", 10.0, 1.5, True, 5.0), Candidate("B", 20.0, 1.2, True, 3.0), Candidate("C", 5.0, 1.0, False, 0.0), Candidate("D", 2.0, 0.5, False, 0.0), Candidate("E", 50.0, 0.9, True, 8.0), ] def test_allocates_three_buckets(self): res = allocate_capital(200_000, self.make_candidates()) bucket_notional = res.bucket_notional self.assertIn(1, bucket_notional) self.assertIn(2, bucket_notional) self.assertIn(3, bucket_notional) # total invested + unallocated == capital self.assertAlmostEqual(res.invested + res.unallocated_cash, 200_000, places=2) # every order is a multiple of 100 for o in res.orders: self.assertEqual(o.qty % 100, 0) self.assertGreaterEqual(o.qty, 100) def test_min_100_shares_respected(self): # tiny capital: bucket1 50% still must afford 100 shares res = allocate_capital(2_000, self.make_candidates()) for o in res.orders: self.assertGreaterEqual(o.qty, 100) self.assertEqual(o.notional, o.qty * o.price) # never over-invest beyond capital self.assertLessEqual(res.invested, 2_000) def test_bucket3_excludes_already_bought(self): cands = [Candidate("A", 10.0, 1.5, True, 5.0), Candidate("B", 10.0, 0.1, True, 8.0)] res = allocate_capital(300_000, cands) # A (bought in bucket1) must NOT also appear in bucket3 bucket3_syms = [o.symbol for o in res.orders if o.bucket == 3] bucket1_syms = [o.symbol for o in res.orders if o.bucket == 1] overlap = set(bucket3_syms) & set(bucket1_syms) self.assertEqual(overlap, set()) def test_bucket3_ranks_by_yield_ignoring_score(self): # The spec: bucket3 = highest dividend yield, IGNORING score. # A low-score but high-yield name must rank above a high-score low-yield name. cands = [ Candidate("HIGH_SCORE", 10.0, 8.0, True, 2.0), # score 8, yield 2% Candidate("HIGH_YIELD", 10.0, 0.1, True, 7.0), # score 0.1, yield 7% Candidate("MID", 10.0, 5.0, True, 3.0), ] # Big capital so bucket1 consumes only the top score name, leaving # HIGH_YIELD (not HIGH_SCORE) to be the bucket3 top pick. res = allocate_capital(1_000_000, cands) b3 = [o.symbol for o in res.orders if o.bucket == 3] # HIGH_YIELD (7%) should be selected in bucket3 before HIGH_SCORE (2%) self.assertIn("HIGH_YIELD", b3) if "HIGH_SCORE" in b3: hi = b3.index("HIGH_SCORE") hy = b3.index("HIGH_YIELD") self.assertLess(hy, hi) def test_invalid_capital_raises(self): with self.assertRaises(Exception): allocate_capital(0, self.make_candidates()) with self.assertRaises(Exception): allocate_capital(-100, self.make_candidates()) def test_no_candidates_raises(self): with self.assertRaises(Exception): allocate_capital(100_000, []) def test_cash_fallback_when_nothing_fits(self): # only very high-priced names, tiny capital -> cannot buy 100 shares -> cash cands = [Candidate("X", 1000.0, 2.0, True, 3.0)] res = allocate_capital(50_000, cands) # 50% = 25k, can't buy 100*1000 self.assertEqual(res.orders, []) self.assertAlmostEqual(res.unallocated_cash, 50_000, places=2) class MomentumSelectionTest(unittest.TestCase): """Owner's "ทำกำไร" rule: bucket 1/2 rank by price-trend momentum and are gated on a POSITIVE theme signal — NOT by combined/EPS score.""" def _cand(self, sym, price, combined, div, yield_, momentum, theme): return Candidate(sym, price, combined, div, yield_, momentum=momentum, theme_signal=theme) def test_bucket1_ranks_by_momentum_not_score(self): # A high-score but NEGATIVE-momentum dividend name must NOT win bucket 1 # over a positive-momentum one (owner: "ทำกำไร" = price likely to rise). cands = [ self._cand("SLOW", 10.0, 9.0, True, 2.0, -1.5, 0.5), # high score, falling price self._cand("FAST", 10.0, 0.5, True, 1.0, +2.0, 0.6), # low score, rising price ] res = allocate_capital(1_000_000, cands) b1 = [o.symbol for o in res.orders if o.bucket == 1] self.assertEqual(b1, ["FAST"]) # FAST (momentum +2) beats SLOW (-1.5) def test_negative_theme_is_excluded_from_profit_bucket(self): # Even a high-momentum name is not eligible for bucket 1/2 when its # theme_signal is <= 0 (owner gate). It may still land in bucket 3 (yield). cands = [ self._cand("NO_THEME", 10.0, 5.0, True, 8.0, +3.0, -0.2), # momentum up but theme negative self._cand("GOOD", 10.0, 5.0, True, 2.0, +1.0, 0.5), ] res = allocate_capital(1_000_000, cands) b1 = [o.symbol for o in res.orders if o.bucket == 1] self.assertNotIn("NO_THEME", b1) # gated out of the profit bucket self.assertIn("GOOD", b1) # NO_THEME may still be picked by bucket 3 (highest yield, 8%) b3 = [o.symbol for o in res.orders if o.bucket == 3] self.assertIn("NO_THEME", b3) def test_bucket2_uses_non_dividend_momentum(self): cands = [ self._cand("ND_UP", 10.0, 9.0, False, 0.0, +2.5, 0.7), # non-div, momentum up self._cand("ND_DN", 10.0, 9.0, False, 0.0, -2.5, 0.7), # non-div, momentum down ] res = allocate_capital(1_000_000, cands) b2 = [o.symbol for o in res.orders if o.bucket == 2] self.assertEqual(b2, ["ND_UP"]) # rising price chosen over falling def test_negative_momentum_excluded_from_profit_buckets(self): # A rising-price name wins the profit bucket; a FALLING-price name — # even with a positive theme — must NOT be picked as "ทำกำไร". cands = [ self._cand("UP", 10.0, 5.0, False, 0.0, +1.5, 0.6), # rising, non-div self._cand("DOWN", 10.0, 5.0, False, 0.0, -1.5, 0.6), # falling, same theme ] res = allocate_capital(1_000_000, cands) b2 = [o.symbol for o in res.orders if o.bucket == 2] self.assertEqual(b2, ["UP"]) # DOWN (falling price) excluded self.assertNotIn("DOWN", b2) class PriceSnapshotTest(unittest.TestCase): """Regression: load_price_snapshot must pick the snapshot retrieved MOST RECENTLY, not the last filename lexicographically (a partial 9-symbol file can sort after a full 50-symbol one and silently drop the universe).""" def test_load_picks_most_recently_retrieved(self): import json, tempfile from pathlib import Path from unittest import mock from app import simulation as sim def _snap(symbols, retrieved_at): return { "schema_version": 1, "source": {"retrieved_at": retrieved_at, "period_start": "2024-01-01", "period_end": "2026-08-29"}, "series": {s: {"bars": [{"date": "2026-08-29", "adjusted_close": 10.0}]} for s in symbols}, } with tempfile.TemporaryDirectory() as td: snap_dir = Path(td) / "snapshots" snap_dir.mkdir() # name that sorts LAST lexicographically, but is OLD (partial 9) old_partial = snap_dir / "prices-yahoo-chart-2024-01-01-2026-08-24-zzz.json" old_partial.write_text(json.dumps(_snap(["A", "B"], "2026-08-23T00:00:00+00:00"))) # full 50-symbol, retrieved MORE recently — should win full = snap_dir / "prices-yahoo-chart-2023-08-28-2026-08-29-aaa.json" full.write_text(json.dumps(_snap([f"S{i}" for i in range(50)], "2026-08-30T00:00:00+00:00"))) with mock.patch.object(sim, "_PRICES_DIR", Path(td)): series = sim.load_price_snapshot() self.assertEqual(len(series), 50) # the full universe, not the 9 if __name__ == "__main__": unittest.main()