blackopsrepl/portfolio-optimization-python

1"""
2Tests for business metrics in the Portfolio Optimization quickstart.
3
4These tests verify the financial KPIs calculated by the domain model:
5- Herfindahl-Hirschman Index (HHI) for concentration
6- Diversification score (1 - HHI)
7- Max sector exposure
8- Expected return
9- Return volatility
10- Sharpe proxy (return / volatility)
11
12These metrics provide business insight beyond the solver score.
13"""
14import pytest
15import math
16
17from portfolio_optimization.domain import (
18    StockSelection,
19    PortfolioOptimizationPlan,
20    PortfolioConfig,
21    PortfolioMetricsModel,
22    SELECTED,
23    NOT_SELECTED,
24)
25from portfolio_optimization.converters import plan_to_metrics
26
27
28def create_stock(
29    stock_id: str,
30    sector: str = "Technology",
31    predicted_return: float = 0.10,
32    selected: bool = True
33) -> StockSelection:
34    """Create a test stock with sensible defaults."""
35    return StockSelection(
36        stock_id=stock_id,
37        stock_name=f"{stock_id} Corp",
38        sector=sector,
39        predicted_return=predicted_return,
40        selection=SELECTED if selected else NOT_SELECTED,
41    )
42
43
44def create_plan(stocks: list[StockSelection]) -> PortfolioOptimizationPlan:
45    """Create a test plan with given stocks."""
46    return PortfolioOptimizationPlan(
47        stocks=stocks,
48        target_position_count=20,
49        max_sector_percentage=0.25,
50        portfolio_config=PortfolioConfig(target_count=20, max_per_sector=5, unselected_penalty=10000),
51    )
52
53
54class TestHerfindahlIndex:
55    """Tests for the Herfindahl-Hirschman Index (HHI) calculation."""
56
57    def test_single_sector_hhi_is_one(self) -> None:
58        """All stocks in one sector should have HHI = 1.0 (max concentration)."""
59        stocks = [create_stock(f"STK{i}", sector="Technology") for i in range(5)]
60        plan = create_plan(stocks)
61
62        # All in one sector: HHI = 1.0^2 = 1.0
63        assert plan.get_herfindahl_index() == 1.0
64
65    def test_two_equal_sectors_hhi(self) -> None:
66        """Two sectors with equal stocks should have HHI = 0.5."""
67        stocks = [
68            *[create_stock(f"TECH{i}", sector="Technology") for i in range(5)],
69            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(5)],
70        ]
71        plan = create_plan(stocks)
72
73        # 50% in each sector: HHI = 0.5^2 + 0.5^2 = 0.5
74        assert abs(plan.get_herfindahl_index() - 0.5) < 0.001
75
76    def test_four_equal_sectors_hhi(self) -> None:
77        """Four sectors with equal stocks should have HHI = 0.25."""
78        stocks = [
79            *[create_stock(f"TECH{i}", sector="Technology") for i in range(5)],
80            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(5)],
81            *[create_stock(f"FIN{i}", sector="Finance") for i in range(5)],
82            *[create_stock(f"NRG{i}", sector="Energy") for i in range(5)],
83        ]
84        plan = create_plan(stocks)
85
86        # 25% in each sector: HHI = 4 * 0.25^2 = 0.25
87        assert abs(plan.get_herfindahl_index() - 0.25) < 0.001
88
89    def test_empty_portfolio_hhi_is_zero(self) -> None:
90        """Empty portfolio should have HHI = 0."""
91        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
92        plan = create_plan(stocks)
93
94        assert plan.get_herfindahl_index() == 0.0
95
96    def test_unequal_sectors_hhi(self) -> None:
97        """Unequal sector distribution should give correct HHI."""
98        stocks = [
99            *[create_stock(f"TECH{i}", sector="Technology") for i in range(6)],  # 60%
100            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(4)],  # 40%
101        ]
102        plan = create_plan(stocks)
103
104        # HHI = 0.6^2 + 0.4^2 = 0.36 + 0.16 = 0.52
105        assert abs(plan.get_herfindahl_index() - 0.52) < 0.001
106
107
108class TestDiversificationScore:
109    """Tests for the diversification score (1 - HHI)."""
110
111    def test_single_sector_diversification_is_zero(self) -> None:
112        """All stocks in one sector should have diversification = 0."""
113        stocks = [create_stock(f"STK{i}", sector="Technology") for i in range(5)]
114        plan = create_plan(stocks)
115
116        assert plan.get_diversification_score() == 0.0
117
118    def test_two_equal_sectors_diversification(self) -> None:
119        """Two equal sectors should have diversification = 0.5."""
120        stocks = [
121            *[create_stock(f"TECH{i}", sector="Technology") for i in range(5)],
122            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(5)],
123        ]
124        plan = create_plan(stocks)
125
126        assert abs(plan.get_diversification_score() - 0.5) < 0.001
127
128    def test_four_equal_sectors_diversification(self) -> None:
129        """Four equal sectors should have diversification = 0.75."""
130        stocks = [
131            *[create_stock(f"TECH{i}", sector="Technology") for i in range(5)],
132            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(5)],
133            *[create_stock(f"FIN{i}", sector="Finance") for i in range(5)],
134            *[create_stock(f"NRG{i}", sector="Energy") for i in range(5)],
135        ]
136        plan = create_plan(stocks)
137
138        # 1 - HHI = 1 - 0.25 = 0.75
139        assert abs(plan.get_diversification_score() - 0.75) < 0.001
140
141
142class TestMaxSectorExposure:
143    """Tests for max sector exposure calculation."""
144
145    def test_single_sector_max_exposure_is_one(self) -> None:
146        """All stocks in one sector should have max exposure = 1.0."""
147        stocks = [create_stock(f"STK{i}", sector="Technology") for i in range(5)]
148        plan = create_plan(stocks)
149
150        assert plan.get_max_sector_exposure() == 1.0
151
152    def test_two_equal_sectors_max_exposure(self) -> None:
153        """Two equal sectors should have max exposure = 0.5."""
154        stocks = [
155            *[create_stock(f"TECH{i}", sector="Technology") for i in range(5)],
156            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(5)],
157        ]
158        plan = create_plan(stocks)
159
160        assert abs(plan.get_max_sector_exposure() - 0.5) < 0.001
161
162    def test_unequal_sectors_max_exposure(self) -> None:
163        """Unequal sectors should return the larger weight."""
164        stocks = [
165            *[create_stock(f"TECH{i}", sector="Technology") for i in range(7)],  # 70%
166            *[create_stock(f"HLTH{i}", sector="Healthcare") for i in range(3)],  # 30%
167        ]
168        plan = create_plan(stocks)
169
170        assert abs(plan.get_max_sector_exposure() - 0.7) < 0.001
171
172    def test_empty_portfolio_max_exposure_is_zero(self) -> None:
173        """Empty portfolio should have max exposure = 0."""
174        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
175        plan = create_plan(stocks)
176
177        assert plan.get_max_sector_exposure() == 0.0
178
179
180class TestSectorCount:
181    """Tests for sector count calculation."""
182
183    def test_single_sector(self) -> None:
184        """All stocks in one sector should return count = 1."""
185        stocks = [create_stock(f"STK{i}", sector="Technology") for i in range(5)]
186        plan = create_plan(stocks)
187
188        assert plan.get_sector_count() == 1
189
190    def test_multiple_sectors(self) -> None:
191        """Stocks in multiple sectors should return correct count."""
192        stocks = [
193            create_stock("TECH1", sector="Technology"),
194            create_stock("HLTH1", sector="Healthcare"),
195            create_stock("FIN1", sector="Finance"),
196            create_stock("NRG1", sector="Energy"),
197        ]
198        plan = create_plan(stocks)
199
200        assert plan.get_sector_count() == 4
201
202    def test_empty_portfolio_sector_count_is_zero(self) -> None:
203        """Empty portfolio should have sector count = 0."""
204        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
205        plan = create_plan(stocks)
206
207        assert plan.get_sector_count() == 0
208
209
210class TestExpectedReturn:
211    """Tests for expected return calculation."""
212
213    def test_uniform_returns(self) -> None:
214        """Stocks with same returns should give that return."""
215        stocks = [create_stock(f"STK{i}", predicted_return=0.10) for i in range(5)]
216        plan = create_plan(stocks)
217
218        assert abs(plan.get_expected_return() - 0.10) < 0.001
219
220    def test_mixed_returns(self) -> None:
221        """Mixed returns should give weighted average."""
222        stocks = [
223            create_stock("STK1", predicted_return=0.10),  # 10%
224            create_stock("STK2", predicted_return=0.20),  # 20%
225        ]
226        plan = create_plan(stocks)
227
228        # Equal weight: (0.10 + 0.20) / 2 = 0.15
229        assert abs(plan.get_expected_return() - 0.15) < 0.001
230
231    def test_empty_portfolio_return_is_zero(self) -> None:
232        """Empty portfolio should have return = 0."""
233        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
234        plan = create_plan(stocks)
235
236        assert plan.get_expected_return() == 0.0
237
238
239class TestReturnVolatility:
240    """Tests for return volatility (std dev) calculation."""
241
242    def test_uniform_returns_zero_volatility(self) -> None:
243        """All same returns should give volatility = 0."""
244        stocks = [create_stock(f"STK{i}", predicted_return=0.10) for i in range(5)]
245        plan = create_plan(stocks)
246
247        assert plan.get_return_volatility() == 0.0
248
249    def test_varied_returns_nonzero_volatility(self) -> None:
250        """Varied returns should give positive volatility."""
251        stocks = [
252            create_stock("STK1", predicted_return=0.05),
253            create_stock("STK2", predicted_return=0.10),
254            create_stock("STK3", predicted_return=0.15),
255            create_stock("STK4", predicted_return=0.20),
256        ]
257        plan = create_plan(stocks)
258
259        # Mean = 0.125, variance = ((0.05-0.125)^2 + (0.10-0.125)^2 + (0.15-0.125)^2 + (0.20-0.125)^2) / 4
260        # = (0.005625 + 0.000625 + 0.000625 + 0.005625) / 4 = 0.003125
261        # Std dev = sqrt(0.003125) ≈ 0.0559
262        expected_vol = math.sqrt(0.003125)
263        assert abs(plan.get_return_volatility() - expected_vol) < 0.0001
264
265    def test_single_stock_zero_volatility(self) -> None:
266        """Single stock should have volatility = 0 (need at least 2)."""
267        stocks = [create_stock("STK1", predicted_return=0.10)]
268        plan = create_plan(stocks)
269
270        assert plan.get_return_volatility() == 0.0
271
272
273class TestSharpeProxy:
274    """Tests for Sharpe ratio proxy calculation."""
275
276    def test_positive_sharpe(self) -> None:
277        """Positive return with volatility should give positive Sharpe."""
278        stocks = [
279            create_stock("STK1", predicted_return=0.05),
280            create_stock("STK2", predicted_return=0.10),
281            create_stock("STK3", predicted_return=0.15),
282            create_stock("STK4", predicted_return=0.20),
283        ]
284        plan = create_plan(stocks)
285
286        # Return = 0.125, volatility = 0.0559
287        # Sharpe = 0.125 / 0.0559 ≈ 2.24
288        sharpe = plan.get_sharpe_proxy()
289        assert sharpe > 2.0
290        assert sharpe < 2.5
291
292    def test_zero_volatility_zero_sharpe(self) -> None:
293        """Zero volatility should give Sharpe = 0 (undefined)."""
294        stocks = [create_stock(f"STK{i}", predicted_return=0.10) for i in range(5)]
295        plan = create_plan(stocks)
296
297        assert plan.get_sharpe_proxy() == 0.0
298
299    def test_empty_portfolio_zero_sharpe(self) -> None:
300        """Empty portfolio should have Sharpe = 0."""
301        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
302        plan = create_plan(stocks)
303
304        assert plan.get_sharpe_proxy() == 0.0
305
306
307class TestPlanToMetrics:
308    """Tests for the plan_to_metrics converter function."""
309
310    def test_metrics_from_valid_portfolio(self) -> None:
311        """plan_to_metrics should return all metrics for valid portfolio."""
312        stocks = [
313            *[create_stock(f"TECH{i}", sector="Technology", predicted_return=0.12) for i in range(5)],
314            *[create_stock(f"HLTH{i}", sector="Healthcare", predicted_return=0.08) for i in range(5)],
315        ]
316        plan = create_plan(stocks)
317
318        metrics = plan_to_metrics(plan)
319
320        assert metrics is not None
321        assert isinstance(metrics, PortfolioMetricsModel)
322        assert metrics.sector_count == 2
323        assert abs(metrics.expected_return - 0.10) < 0.001
324        assert abs(metrics.diversification_score - 0.5) < 0.001
325        assert abs(metrics.herfindahl_index - 0.5) < 0.001
326        assert abs(metrics.max_sector_exposure - 0.5) < 0.001
327
328    def test_metrics_from_empty_portfolio_is_none(self) -> None:
329        """plan_to_metrics should return None for empty portfolio."""
330        stocks = [create_stock(f"STK{i}", selected=False) for i in range(5)]
331        plan = create_plan(stocks)
332
333        metrics = plan_to_metrics(plan)
334
335        assert metrics is None
336
337    def test_metrics_serialization(self) -> None:
338        """Metrics should serialize with camelCase aliases."""
339        stocks = [create_stock(f"STK{i}") for i in range(5)]
340        plan = create_plan(stocks)
341
342        metrics = plan_to_metrics(plan)
343        assert metrics is not None
344
345        data = metrics.model_dump(by_alias=True)
346        assert "expectedReturn" in data
347        assert "sectorCount" in data
348        assert "maxSectorExposure" in data
349        assert "herfindahlIndex" in data
350        assert "diversificationScore" in data
351        assert "returnVolatility" in data
352        assert "sharpeProxy" in data
353