From de77d5b2aa9db7a30dc2785eb056cc028bf66938 Mon Sep 17 00:00:00 2001
From: zephyrdark <zephyrdark@gmail.com>
Date: Wed, 11 Feb 2026 23:27:47 +0900
Subject: [PATCH] feat: add rebalance calculate schemas and tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 backend/app/schemas/portfolio.py         |  32 +++++++
 backend/tests/e2e/test_rebalance_flow.py | 113 +++++++++++++++++++++++
 2 files changed, 145 insertions(+)
 create mode 100644 backend/tests/e2e/test_rebalance_flow.py

diff --git a/backend/app/schemas/portfolio.py b/backend/app/schemas/portfolio.py
index db6a4d3..5ae8eb4 100644
--- a/backend/app/schemas/portfolio.py
+++ b/backend/app/schemas/portfolio.py
@@ -186,3 +186,35 @@ class RebalanceSimulationResponse(BaseModel):
     additional_amount: Decimal
     new_total: Decimal
     items: List[RebalanceItem]
+
+
+class RebalanceCalculateRequest(BaseModel):
+    """Request for manual rebalance calculation."""
+    strategy: str = Field(..., pattern="^(full_rebalance|additional_buy)$")
+    prices: Optional[dict[str, Decimal]] = None
+    additional_amount: Optional[Decimal] = Field(None, ge=0)
+
+
+class RebalanceCalculateItem(BaseModel):
+    """Extended rebalance item with price change info."""
+    ticker: str
+    name: str | None = None
+    target_ratio: Decimal
+    current_ratio: Decimal
+    current_quantity: int
+    current_value: Decimal
+    current_price: Decimal
+    target_value: Decimal
+    diff_ratio: Decimal
+    diff_quantity: int
+    action: str  # "buy", "sell", or "hold"
+    change_vs_prev_month: Decimal | None = None
+    change_vs_start: Decimal | None = None
+
+
+class RebalanceCalculateResponse(BaseModel):
+    """Response for manual rebalance calculation."""
+    portfolio_id: int
+    total_assets: Decimal
+    available_to_buy: Decimal | None = None
+    items: List[RebalanceCalculateItem]
diff --git a/backend/tests/e2e/test_rebalance_flow.py b/backend/tests/e2e/test_rebalance_flow.py
new file mode 100644
index 0000000..95ad9de
--- /dev/null
+++ b/backend/tests/e2e/test_rebalance_flow.py
@@ -0,0 +1,113 @@
+"""
+E2E tests for rebalancing calculation flow.
+"""
+from fastapi.testclient import TestClient
+
+
+def _setup_portfolio_with_holdings(client: TestClient, auth_headers: dict) -> int:
+    """Helper: create portfolio with targets and holdings."""
+    # Create portfolio
+    resp = client.post(
+        "/api/portfolios",
+        json={"name": "Rebalance Test", "portfolio_type": "pension"},
+        headers=auth_headers,
+    )
+    pid = resp.json()["id"]
+
+    # Set targets (sum = 100)
+    client.put(
+        f"/api/portfolios/{pid}/targets",
+        json=[
+            {"ticker": "069500", "target_ratio": 50},
+            {"ticker": "148070", "target_ratio": 50},
+        ],
+        headers=auth_headers,
+    )
+
+    # Set holdings
+    client.put(
+        f"/api/portfolios/{pid}/holdings",
+        json=[
+            {"ticker": "069500", "quantity": 10, "avg_price": 40000},
+            {"ticker": "148070", "quantity": 5, "avg_price": 100000},
+        ],
+        headers=auth_headers,
+    )
+    return pid
+
+
+def test_calculate_rebalance_with_manual_prices(client: TestClient, auth_headers):
+    """Test rebalance calculation with manually provided prices."""
+    pid = _setup_portfolio_with_holdings(client, auth_headers)
+
+    response = client.post(
+        f"/api/portfolios/{pid}/rebalance/calculate",
+        json={
+            "strategy": "full_rebalance",
+            "prices": {"069500": 50000, "148070": 110000},
+        },
+        headers=auth_headers,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["portfolio_id"] == pid
+    assert data["total_assets"] > 0
+    assert len(data["items"]) == 2
+    # Verify items have required fields
+    item = data["items"][0]
+    assert "ticker" in item
+    assert "target_ratio" in item
+    assert "current_ratio" in item
+    assert "diff_quantity" in item
+    assert "action" in item
+    assert "change_vs_prev_month" in item
+
+
+def test_calculate_additional_buy_strategy(client: TestClient, auth_headers):
+    """Test additional buy strategy: buy-only, no sells."""
+    pid = _setup_portfolio_with_holdings(client, auth_headers)
+
+    response = client.post(
+        f"/api/portfolios/{pid}/rebalance/calculate",
+        json={
+            "strategy": "additional_buy",
+            "prices": {"069500": 50000, "148070": 110000},
+            "additional_amount": 1000000,
+        },
+        headers=auth_headers,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["total_assets"] > 0
+    assert data["available_to_buy"] == 1000000
+    # Additional buy should never have "sell" actions
+    for item in data["items"]:
+        assert item["action"] in ("buy", "hold")
+
+
+def test_additional_buy_requires_amount(client: TestClient, auth_headers):
+    """Test that additional_buy strategy requires additional_amount."""
+    pid = _setup_portfolio_with_holdings(client, auth_headers)
+
+    response = client.post(
+        f"/api/portfolios/{pid}/rebalance/calculate",
+        json={
+            "strategy": "additional_buy",
+            "prices": {"069500": 50000, "148070": 110000},
+        },
+        headers=auth_headers,
+    )
+    assert response.status_code == 400
+
+
+def test_calculate_rebalance_without_prices_fallback(client: TestClient, auth_headers):
+    """Test rebalance calculation without manual prices falls back to DB."""
+    pid = _setup_portfolio_with_holdings(client, auth_headers)
+
+    # Without prices, should still work (may have 0 prices from DB in test env)
+    response = client.post(
+        f"/api/portfolios/{pid}/rebalance/calculate",
+        json={"strategy": "full_rebalance"},
+        headers=auth_headers,
+    )
+    assert response.status_code == 200