diff --git a/backend/app/simulation.py b/backend/app/simulation.py index 74be6e0..97d0da4 100644 --- a/backend/app/simulation.py +++ b/backend/app/simulation.py @@ -133,10 +133,16 @@ def allocate_capital( used = set() - def _fill(bucket_idx: int, eligible: list[Candidate], require_dividend: bool): + def _fill(bucket_idx: int, eligible: list[Candidate], require_dividend: bool, + sort_by: str = "combined_score"): nonlocal cash, used remaining = cash[bucket_idx] - for cand in sorted(eligible, key=lambda c: -c.combined_score): + # bucket 3 must rank by dividend_yield (ignoring score); others by score. + if sort_by == "dividend_yield": + key = lambda c: -c.dividend_yield + else: + key = lambda c: -c.combined_score + for cand in sorted(eligible, key=key): if cand.symbol in used: continue if require_dividend and not cand.is_dividend: @@ -164,8 +170,8 @@ def allocate_capital( _fill(0, [c for c in by_score if c.is_dividend], require_dividend=True) # Bucket 2: non-dividend, highest score _fill(1, [c for c in by_score if not c.is_dividend], require_dividend=False) - # Bucket 3: highest dividend yield, excluding symbols already bought - _fill(2, by_yield, require_dividend=True) + # Bucket 3: highest dividend yield (ignoring score), excluding symbols bought + _fill(2, by_yield, require_dividend=True, sort_by="dividend_yield") result.unallocated_cash = sum(cash) return result diff --git a/backend/tests/test_simulation.py b/backend/tests/test_simulation.py index 8614dd7..71ca943 100644 --- a/backend/tests/test_simulation.py +++ b/backend/tests/test_simulation.py @@ -50,6 +50,25 @@ class AllocationTest(unittest.TestCase): 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())