diff --git a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py index 5cd8bb7cc..2a75c51af 100644 --- a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py +++ b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py @@ -133,7 +133,7 @@ class GeodataFileTable(RootModel): root: list[GeodataFileRow] @model_validator(mode="after") - def subpop_is_primary_key(self) -> "GeodataFileTable": + def _subpop_is_primary_key(self) -> "GeodataFileTable": """ Validate that the subpopulation names are unique. @@ -163,10 +163,10 @@ class MobilityFileRow(BaseModel): ori: str dest: str - amount: Annotated[int, Field(gt=0)] + amount: Annotated[int, Field(ge=0)] @model_validator(mode="after") - def ori_and_dest_are_different(self) -> "MobilityFileRow": + def _ori_and_dest_are_different(self) -> "MobilityFileRow": """ Validate that the origin and destination subpopulation names are different. @@ -194,7 +194,7 @@ class MobilityFileTable(RootModel): root: list[MobilityFileRow] @model_validator(mode="after") - def ori_and_dest_are_primary_key(self) -> "MobilityFileTable": + def _ori_and_dest_are_primary_key(self) -> "MobilityFileTable": """ Validate that the origin and destination subpopulation names are unique. diff --git a/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py b/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py index 7ae75760a..cfbad73a9 100644 --- a/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py +++ b/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py @@ -365,6 +365,51 @@ def valid_2pop_with_csv_mobility_factory( ) +def valid_2pop_with_csv_explicit_no_mobility_factory( + tmp_path: Path, +) -> MockSubpopulationStructureInput: + """ + Factory for geodata file with two valid subpopulations and explicit no mobility. + + Returns: + A `MockSubpopulationStructureInput` instance with a geodata file that contains + two valid subpopulations and a mobility matrix. + """ + return MockSubpopulationStructureInput.create_mock_input( + tmp_path, + { + "geodata": "geodata.csv", + "mobility": "mobility.csv", + }, + pd.DataFrame.from_records( + [ + { + "subpop": "USA", + "population": 100, + }, + { + "subpop": "Canada", + "population": 50, + }, + ] + ), + mobility=pd.DataFrame.from_records( + [ + { + "ori": "USA", + "dest": "Canada", + "amount": 0, + }, + { + "ori": "Canada", + "dest": "USA", + "amount": 2, + }, + ] + ), + ) + + def valid_2pop_with_parquet_mobility_factory( tmp_path: Path, ) -> MockSubpopulationStructureInput: @@ -684,6 +729,29 @@ def mobility_zero_or_less_factory(tmp_path: Path) -> MockSubpopulationStructureI ) +def larger_mobility_zero_or_less_factory(tmp_path: Path) -> MockSubpopulationStructureInput: + return MockSubpopulationStructureInput.create_mock_input( + tmp_path, + { + "geodata": "geodata.parquet", + "mobility": "mobility.parquet", + }, + pd.DataFrame( + data={ + "subpop": ["USA", "Canada", "Mexico"], + "population": [100, 50, 25], + }, + ), + mobility=pd.DataFrame( + data={ + "ori": ["USA", "Canada", "Mexico", "Canada", "USA", "Mexico"], + "dest": ["Canada", "USA", "Canada", "Mexico", "Mexico", "USA"], + "amount": [0, -1, 1, 0, -1, 0], + }, + ), + ) + + def selected_missing_from_geodata(tmp_path: Path) -> MockSubpopulationStructureInput: return MockSubpopulationStructureInput.create_mock_input( tmp_path, @@ -711,6 +779,7 @@ def selected_missing_from_geodata(tmp_path: Path) -> MockSubpopulationStructureI valid_2pop_geodata_parquet_only_factory, valid_2pop_with_txt_mobility_factory, valid_2pop_with_csv_mobility_factory, + valid_2pop_with_csv_explicit_no_mobility_factory, valid_2pop_with_parquet_mobility_factory, valid_2pop_with_npz_mobility_factory, valid_selected_pop_with_no_mobility_factory, @@ -810,13 +879,25 @@ def test_subpopulation_structure_instance_attributes( assert len(caplog.records) == int(mock_input.mobility is None) -def test_mobility_zero_or_less_raises_validation_error(tmp_path: Path) -> None: - """Test that a ValueError is raised when mobility is zero or less.""" - mock_input = mobility_zero_or_less_factory(tmp_path) - assert (zero_or_less := (mock_input.mobility["amount"] <= 0).sum()) > 0 - raises_match = rf"^{zero_or_less} validation errors.*" +@pytest.mark.parametrize( + "factory", [mobility_zero_or_less_factory, larger_mobility_zero_or_less_factory] +) +def test_mobility_less_than_zero_raises_validation_error( + tmp_path: Path, factory: Callable[[Path], MockSubpopulationStructureInput] +) -> None: + """Test that a ValueError is raised when mobility is less than zero.""" + mock_input = factory(tmp_path) + assert ( + less_than_zero_indices := ( + mock_input.mobility.index[mock_input.mobility["amount"] < 0] + ).tolist() + ) + raises_match = rf"^{len(less_than_zero_indices)} validation error.*" raises_match += ".*".join( - [rf"{i}\.amount\s+Input should be greater than 0" for i in range(zero_or_less)] + [ + rf"{i}\.amount\s+Input should be greater than or equal to 0" + for i in less_than_zero_indices + ] ) with pytest.raises( ValidationError,