Allow fully plural stored household groups

Author: anth-volk

src/policyengine_api/models/household.py

@@ -20,7 +20,6 @@
 }
 
 _ENTITY_GROUP_FIELDS = tuple(_ENTITY_ID_KEY_BY_GROUP.keys())
-_MULTI_ROW_ALLOWED = {"marital_unit"}
 
 
 def _coerce_entity_group_collection(value: Any) -> Any:
@@ -55,8 +54,8 @@ class Household(HouseholdBase, table=True):
 class HouseholdCreate(StoredHouseholdPayload):
     """Schema for creating a stored household.
 
-    Uses list-based entity groups for storage, with multiple rows allowed only
-    for marital_unit.
+    Uses the same list-based relational entity shape as the household
+    calculation APIs.
     """
 
     @field_validator(*_ENTITY_GROUP_FIELDS, mode="before")
@@ -78,11 +77,6 @@ def validate_relationships(self) -> "HouseholdCreate":
             entity_records = getattr(self, group_key)
             person_link_key = f"person_{entity_id_key}"
 
-            if len(entity_records) > 1 and group_key not in _MULTI_ROW_ALLOWED:
-                raise ValueError(
-                    f"{group_key} supports at most one row in stored households"
-                )
-
             entity_ids = [
                 entity[entity_id_key]
                 for entity in entity_records

src/policyengine_api/models/household_payload.py

@@ -8,7 +8,7 @@
 
 
 class HouseholdEntityCollections(SQLModel):
-    """Plural household entity collections used by calculation and storage APIs."""
+    """Plural household entity collections used by stored and calculation APIs."""
 
     benunit: list[dict[str, Any]] = Field(default_factory=list)
     marital_unit: list[dict[str, Any]] = Field(default_factory=list)

test_fixtures/fixtures_households.py

@@ -62,6 +62,54 @@
     "household": [{"household_id": 0, "state_name": "CA"}],
 }
 
+MOCK_US_FULL_MULTI_GROUP_HOUSEHOLD_CREATE = {
+    "country_id": "us",
+    "year": 2024,
+    "label": "US fully multi-group household",
+    "people": [
+        {
+            "person_id": 0,
+            "person_household_id": 0,
+            "person_tax_unit_id": 0,
+            "person_marital_unit_id": 0,
+            "person_family_id": 0,
+            "person_spm_unit_id": 0,
+            "age": 30,
+            "employment_income": 50000,
+        },
+        {
+            "person_id": 1,
+            "person_household_id": 1,
+            "person_tax_unit_id": 1,
+            "person_marital_unit_id": 1,
+            "person_family_id": 1,
+            "person_spm_unit_id": 1,
+            "age": 28,
+            "employment_income": 30000,
+        },
+    ],
+    "tax_unit": [
+        {"tax_unit_id": 0, "state_name": "CA"},
+        {"tax_unit_id": 1, "state_name": "CA"},
+    ],
+    "marital_unit": [
+        {"marital_unit_id": 0},
+        {"marital_unit_id": 1},
+    ],
+    "family": [
+        {"family_id": 0},
+        {"family_id": 1},
+    ],
+    "spm_unit": [
+        {"spm_unit_id": 0},
+        {"spm_unit_id": 1},
+    ],
+    "household": [
+        {"household_id": 0, "state_name": "CA"},
+        {"household_id": 1, "state_name": "NY"},
+    ],
+}
+
 MOCK_HOUSEHOLD_MINIMAL = {
     "country_id": "us",
     "year": 2024,

tests/test_households.py

@@ -5,6 +5,7 @@
 from test_fixtures.fixtures_households import (
     MOCK_HOUSEHOLD_MINIMAL,
     MOCK_UK_HOUSEHOLD_CREATE,
+    MOCK_US_FULL_MULTI_GROUP_HOUSEHOLD_CREATE,
     MOCK_US_HOUSEHOLD_CREATE,
     MOCK_US_HOUSEHOLD_CREATE_LEGACY,
     MOCK_US_MULTI_GROUP_HOUSEHOLD_CREATE,
@@ -62,14 +63,20 @@ def test_create_household_minimal(client):
 
 
 def test_create_household_round_trips_multiple_entity_groups(client):
-    """Stored household CRUD preserves multiple marital units."""
-    response = client.post("/households", json=MOCK_US_MULTI_GROUP_HOUSEHOLD_CREATE)
+    """Stored household CRUD preserves multiple rows across entity groups."""
+    response = client.post(
+        "/households", json=MOCK_US_FULL_MULTI_GROUP_HOUSEHOLD_CREATE
+    )
     assert response.status_code == 201
     data = response.json()
-    assert len(data["tax_unit"]) == 1
-    assert data["tax_unit"][0]["tax_unit_id"] == 0
+    assert len(data["tax_unit"]) == 2
     assert len(data["marital_unit"]) == 2
-    assert data["people"][1]["person_marital_unit_id"] == 1
+    assert len(data["family"]) == 2
+    assert len(data["spm_unit"]) == 2
+    assert len(data["household"]) == 2
+    assert data["people"][0]["person_tax_unit_id"] == 0
+    assert data["people"][1]["person_tax_unit_id"] == 1
+    assert data["people"][1]["person_household_id"] == 1
 
 
 def test_create_household_invalid_country_id(client):
@@ -154,8 +161,8 @@ def test_create_household_requires_person_links_for_multi_group_rows(client):
     assert "when " in response.text
 
 
-def test_create_household_rejects_multiple_tax_units(client):
-    """Stored households support at most one tax unit."""
+def test_create_household_accepts_multiple_tax_units(client):
+    """Stored households preserve multiple tax units when person links are present."""
     payload = {
         **MOCK_US_MULTI_GROUP_HOUSEHOLD_CREATE,
         "tax_unit": [
@@ -182,12 +189,14 @@ def test_create_household_rejects_multiple_tax_units(client):
 
     response = client.post("/households", json=payload)
 
-    assert response.status_code == 422
-    assert "tax_unit supports at most one row" in response.text
+    assert response.status_code == 201
+    data = response.json()
+    assert len(data["tax_unit"]) == 2
+    assert data["people"][1]["person_tax_unit_id"] == 1
 
 
-def test_create_household_rejects_multiple_households(client):
-    """Stored households support at most one household row."""
+def test_create_household_accepts_multiple_households(client):
+    """Stored households preserve multiple household rows when person links are present."""
     payload = {
         **MOCK_US_MULTI_GROUP_HOUSEHOLD_CREATE,
         "household": [
@@ -214,8 +223,10 @@ def test_create_household_rejects_multiple_households(client):
 
     response = client.post("/households", json=payload)
 
-    assert response.status_code == 422
-    assert "household supports at most one row" in response.text
+    assert response.status_code == 201
+    data = response.json()
+    assert len(data["household"]) == 2
+    assert data["people"][1]["person_household_id"] == 1
 
 
 # ---------------------------------------------------------------------------