mypy fixes in artfiact/test_manager.py (#577)

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**3.2.18 - 01/28/25**
+
+  - Type-hinting: Fix mypy errors in tests/framework/artifact/test_manager.py
+
 **3.2.17 - 01/24/25**
 
   - Type-hinting: Fix mypy errors in tests/framework/randomness/test_index_map.py

pyproject.toml

@@ -43,7 +43,6 @@ exclude = [
     'src/vivarium/testing_utilities.py',
     'tests/examples/test_disease_model.py',
     'tests/framework/artifact/test_hdf.py',
-    'tests/framework/artifact/test_manager.py',
     'tests/framework/components/mocks.py',
     'tests/framework/components/test_component.py',
     'tests/framework/components/test_manager.py',

tests/framework/artifact/test_manager.py

@@ -2,9 +2,12 @@
 
 import random
 from pathlib import Path
+from unittest.mock import MagicMock
 
 import pandas as pd
 import pytest
+from layered_config_tree import LayeredConfigTree
+from pytest_mock import MockerFixture
 
 from vivarium.framework.artifact.manager import (
     ArtifactManager,
@@ -18,10 +21,10 @@
 
 
 @pytest.fixture()
-def artifact_mock(mocker):
+def artifact_mock(mocker: MockerFixture) -> MagicMock:
     mock = mocker.patch("vivarium.framework.artifact.manager.Artifact")
 
-    def mock_load(key):
+    def mock_load(key: str) -> str | pd.DataFrame | None:
         if key == "string_data.key":
             return "string_data"
         elif key == "df_data.key":
@@ -34,13 +37,13 @@ def mock_load(key):
     return mock
 
 
-def test_subset_rows_extra_filters():
+def test_subset_rows_extra_filters() -> None:
     data = build_table(1)
     with pytest.raises(ValueError):
         _subset_rows(data, missing_thing=12)
 
 
-def test_subset_rows():
+def test_subset_rows() -> None:
     values = [
         lambda *args, **kwargs: random.choice(["red", "blue"]),
         lambda *args, **kwargs: random.choice([1, 2, 3]),
@@ -58,7 +61,7 @@ def test_subset_rows():
     )
 
 
-def test_subset_columns():
+def test_subset_columns() -> None:
     values = [0, "red", 100]
     data = build_table(
         values,
@@ -80,7 +83,9 @@ def test_subset_columns():
     )
 
 
-def test_parse_artifact_path_config(base_config, test_data_dir):
+def test_parse_artifact_path_config(
+    base_config: LayeredConfigTree, test_data_dir: Path
+) -> None:
     artifact_path = test_data_dir / "artifact.hdf"
     base_config.update(
         {"input_data": {"artifact_path": str(artifact_path)}}, **metadata(str(Path("/")))
@@ -89,23 +94,25 @@ def test_parse_artifact_path_config(base_config, test_data_dir):
     assert parse_artifact_path_config(base_config) == str(artifact_path)
 
 
-def test_parse_artifact_path_relative_no_source(base_config):
+def test_parse_artifact_path_relative_no_source(base_config: LayeredConfigTree) -> None:
     artifact_path = "./artifact.hdf"
     base_config.update({"input_data": {"artifact_path": str(artifact_path)}})
 
     with pytest.raises(ValueError):
         parse_artifact_path_config(base_config)
 
 
-def test_parse_artifact_path_relative(base_config, test_data_dir):
+def test_parse_artifact_path_relative(
+    base_config: LayeredConfigTree, test_data_dir: Path
+) -> None:
     base_config.update(
         {"input_data": {"artifact_path": "../../test_data/artifact.hdf"}},
         **metadata(__file__),
     )
     assert parse_artifact_path_config(base_config) == str(test_data_dir / "artifact.hdf")
 
 
-def test_parse_artifact_path_config_fail(base_config):
+def test_parse_artifact_path_config_fail(base_config: LayeredConfigTree) -> None:
     artifact_path = Path(__file__).parent / "not_an_artifact.hdf"
     base_config.update(
         {"input_data": {"artifact_path": str(artifact_path)}}, **metadata(str(Path("/")))
@@ -115,7 +122,7 @@ def test_parse_artifact_path_config_fail(base_config):
         parse_artifact_path_config(base_config)
 
 
-def test_parse_artifact_path_config_fail_relative(base_config):
+def test_parse_artifact_path_config_fail_relative(base_config: LayeredConfigTree) -> None:
     base_config.update(
         {"input_data": {"artifact_path": "./not_an_artifact.hdf"}}, **metadata(__file__)
     )
@@ -124,41 +131,41 @@ def test_parse_artifact_path_config_fail_relative(base_config):
         parse_artifact_path_config(base_config)
 
 
-def test_load_with_string_data(artifact_mock):
+def test_load_with_string_data(artifact_mock: MagicMock) -> None:
     am = ArtifactManager()
     am.artifact = artifact_mock
     am.config_filter_term = None
     assert am.load("string_data.key") == "string_data"
 
 
-def test_load_with_no_data(artifact_mock):
+def test_load_with_no_data(artifact_mock: MagicMock) -> None:
     am = ArtifactManager()
     am.artifact = artifact_mock
     assert am.load("no_data.key") is None
 
 
-def test_load_with_df_data(artifact_mock):
+def test_load_with_df_data(artifact_mock: MagicMock) -> None:
     am = ArtifactManager()
     am.artifact = artifact_mock
     am.config_filter_term = None
     assert isinstance(am.load("df_data.key"), pd.DataFrame)
 
 
-def test_config_filter():
+def test_config_filter() -> None:
     df = pd.DataFrame({"year": range(1990, 2000, 1), "color": ["red", "yellow"] * 5})
     filtered = _config_filter(df, "year in [1992, 1995]")
 
     assert set(filtered.year) == {1992, 1995}
 
 
-def test_config_filter_on_nonexistent_column():
+def test_config_filter_on_nonexistent_column() -> None:
     df = pd.DataFrame({"year": range(1990, 2000, 1), "color": ["red", "yellow"] * 5})
     filtered = _config_filter(df, "fake_col in [1992, 1995]")
 
     assert df.equals(filtered)
 
 
-def test_validate_filter_term():
+def test_validate_filter_term() -> None:
     multiple_filter_terms = "draw in [0, 1] and year > 1990"
 
     with pytest.raises(NotImplementedError):