Improve type checking in tests

docs/agents/python.md

@@ -6,12 +6,10 @@ remains configured for editor feedback.
 
 ## Required checks
 
-- **Tool environment:** Run Python checks in this worktree's `.conda-env`:
-  `conda run --prefix .conda-env make typecheck` and
-  `conda run --prefix .conda-env make lint`. If the environment is activated
-  instead, verify that `python`, `mypy`, and `black` all resolve inside the
-  same environment before running the Make targets. A `python` path inside
-  `.conda-env` alone does not establish which tools Make will invoke.
+- **Tool environment:** Run Python checks from within an active, isolated project
+  environment (e.g., `conda`, `venv`, or an agent sandbox). Verify that `python`,
+  `mypy`, and `black` all resolve to this same environment before running the Make
+  targets.
 - **Types:** Annotate functions and methods. Avoid `Any` in signatures. Use
   `make typecheck` for type checking; it runs the project's mypy target.
 - **Format and lint:** Format modified Python files with `black` from the
@@ -27,15 +25,15 @@ remains configured for editor feedback.
 
 ## Testing
 
-- **Test environment:** Run pytest and Make test targets through `.conda-env`.
-  NumCosmo tests write FFTW wisdom under `~/.numcosmo`, and download tests use
-  `~/.firecrown`; ensure these directories are writable before running those
-  tests in a sandbox. Treat a denied cache write as a sandbox failure, then
+- **Test environment:** Run `pytest` and Make test targets within your active
+  environment. NumCosmo tests write FFTW wisdom under `~/.numcosmo`, and download
+  tests use `~/.firecrown`; ensure these directories are writable before running
+  those tests in a sandbox. Treat a denied cache write as a sandbox failure, then
   rerun with scoped access before judging the test result.
-- **Avoid monkeypatching**: Prefer use of fixtures, parameterized tests, and
+- **Avoid monkeypatching:** Prefer use of fixtures, parameterized tests, and
   dependency injection to the use of monkeypatching. Use monkeypatching only
-  when other techniques would lead to more complicated code.  Any use of
-  monkeypatching must contain a comment explaining why monkeypatching it the
+  when other techniques would lead to more complicated code. Any use of
+  monkeypatching must contain a comment explaining why monkeypatching is the
   best choice for that test.
 
 ## Review guidance

pyproject.toml

@@ -148,10 +148,6 @@ explicit_package_bases = true
 disallow_any_generics = true
 disallow_untyped_defs = true
 
-[[tool.mypy.overrides]]
-module = ["tests", "tests.*"]
-disallow_untyped_defs = false
-
 [[tool.mypy.overrides]]
 module = "firecrown.connector.cobaya.*"
 disallow_subclassing_any = false

tests/_types.py

@@ -0,0 +1,36 @@
+"""Shared types for test fixtures."""
+
+from typing import TypeAlias, TypedDict
+
+from numcosmo_py import Nc
+import numpy as np
+import numpy.typing as npt
+import sacc
+
+from firecrown.metadata_types import TracerNames
+
+RedshiftDistribution: TypeAlias = tuple[
+    npt.NDArray[np.float64], npt.NDArray[np.float64]
+]
+"""Sampled redshift distribution as (z, dndz), without a normalization guarantee."""
+
+WindowedAngularPowerSpectrum: TypeAlias = tuple[
+    npt.NDArray[np.int64], npt.NDArray[np.float64], sacc.BandpowerWindow
+]
+"""Window ell grid, measured bandpowers, and their bandpower window."""
+
+SaccGalaxyWindowedFixture: TypeAlias = tuple[
+    sacc.Sacc,
+    dict[str, RedshiftDistribution],
+    dict[tuple[TracerNames, str], WindowedAngularPowerSpectrum],
+]
+"""SACC data, distributions by tracer name, and spectra by tracer pair/data type."""
+
+
+class NumCosmoCosmologyFixture(TypedDict):
+    """Values provided by the NumCosmo cosmology fixtures."""
+
+    cosmo: Nc.HICosmo
+    p_ml: Nc.PowspecML
+    p_mnl: Nc.PowspecMNL
+    dist: Nc.Distance

tests/app/analysis/test_download.py

@@ -72,7 +72,7 @@ class TestDownloadFromUrl:
 
     @patch("urllib.request.urlretrieve")
     def test_download_from_url_uses_cache_when_available(
-        self, mock_retrieve, tmp_path: Path
+        self, mock_retrieve: MagicMock, tmp_path: Path
     ) -> None:
         """Test that cached files are used instead of downloading."""
         output_file = tmp_path / "test.txt"
@@ -99,7 +99,7 @@ def test_download_from_url_uses_cache_when_available(
 
     @patch("urllib.request.urlretrieve")
     def test_download_from_url_downloads_when_not_cached(
-        self, mock_retrieve, tmp_path: Path
+        self, mock_retrieve: MagicMock, tmp_path: Path
     ) -> None:
         """Test that file is downloaded when not in cache."""
         output_file = tmp_path / "test_new.txt"
@@ -113,7 +113,7 @@ def test_download_from_url_downloads_when_not_cached(
 
         try:
             # Mock the download to create the cached file
-            def mock_download(url, path):
+            def mock_download(url: str, path: Path) -> None:
                 assert url is not None
                 Path(path).write_text("downloaded content", encoding="utf-8")
 
@@ -135,7 +135,7 @@ def mock_download(url, path):
 
     @patch("urllib.request.urlretrieve")
     def test_download_from_url_with_description(
-        self, mock_retrieve, tmp_path: Path
+        self, mock_retrieve: MagicMock, tmp_path: Path
     ) -> None:
         """Test download with custom description."""
         output_file = tmp_path / "test_desc.txt"
@@ -149,7 +149,7 @@ def test_download_from_url_with_description(
 
         try:
             # Mock the download
-            def mock_download(url, path):
+            def mock_download(url: str, path: Path) -> None:
                 assert url is not None
                 Path(path).write_text("content", encoding="utf-8")
 
@@ -193,7 +193,7 @@ def test_download_from_url_network_error(self, tmp_path: Path) -> None:
 
     @patch("urllib.request.urlretrieve")
     def test_download_from_url_console_parameter(
-        self, mock_retrieve, tmp_path: Path
+        self, mock_retrieve: MagicMock, tmp_path: Path
     ) -> None:
         """Test that console parameter is required."""
         output_file = tmp_path / "test_console.txt"
@@ -207,7 +207,7 @@ def test_download_from_url_console_parameter(
 
         try:
             # Mock the download
-            def mock_download(_, path):
+            def mock_download(_: str, path: Path) -> None:
                 Path(path).write_text("content", encoding="utf-8")
 
             mock_retrieve.side_effect = mock_download

tests/app/analysis/test_types.py

@@ -3,6 +3,8 @@
 Tests type definitions, enums, and data models for framework configuration.
 """
 
+from pathlib import Path
+
 import pyccl
 import pytest
 from pydantic import ValidationError
@@ -827,7 +829,7 @@ def test_to_ccl_cosmology_all_optional_fields_false(self) -> None:
 class TestConfigGeneratorMethods:
     """Tests for ConfigGenerator methods."""
 
-    def test_config_generator_add_sacc(self, tmp_path) -> None:
+    def test_config_generator_add_sacc(self, tmp_path: Path) -> None:
         """Test add_sacc method."""
         gen = CobayaConfigGenerator(
             output_path=tmp_path,
@@ -841,7 +843,7 @@ def test_config_generator_add_sacc(self, tmp_path) -> None:
         gen.add_sacc(sacc_path)
         assert gen.sacc_path == sacc_path
 
-    def test_config_generator_add_factory(self, tmp_path) -> None:
+    def test_config_generator_add_factory(self, tmp_path: Path) -> None:
         """Test add_factory method."""
         gen = CobayaConfigGenerator(
             output_path=tmp_path,
@@ -855,7 +857,7 @@ def test_config_generator_add_factory(self, tmp_path) -> None:
         gen.add_factory(factory_path)
         assert gen.factory_source == factory_path
 
-    def test_config_generator_add_build_parameters(self, tmp_path) -> None:
+    def test_config_generator_add_build_parameters(self, tmp_path: Path) -> None:
         """Test add_build_parameters method."""
         gen = CobayaConfigGenerator(
             output_path=tmp_path,
@@ -869,7 +871,7 @@ def test_config_generator_add_build_parameters(self, tmp_path) -> None:
         gen.add_build_parameters(params)
         assert gen.build_parameters == params
 
-    def test_config_generator_add_models(self, tmp_path) -> None:
+    def test_config_generator_add_models(self, tmp_path: Path) -> None:
         """Test add_models method."""
         gen = CobayaConfigGenerator(
             output_path=tmp_path,
@@ -908,14 +910,14 @@ def test_get_path_str_with_string(self) -> None:
         result = get_path_str("my_string_path", use_absolute=False)
         assert result == "my_string_path"
 
-    def test_get_path_str_with_path_absolute(self, tmp_path) -> None:
+    def test_get_path_str_with_path_absolute(self, tmp_path: Path) -> None:
         """Test get_path_str with Path object and absolute flag."""
         test_path = tmp_path / "subdir" / "file.txt"
 
         result = get_path_str(test_path, use_absolute=True)
         assert result == test_path.absolute().as_posix()
 
-    def test_get_path_str_with_path_relative(self, tmp_path) -> None:
+    def test_get_path_str_with_path_relative(self, tmp_path: Path) -> None:
         """Test get_path_str with Path object and relative flag."""
         test_path = tmp_path / "subdir" / "file.txt"
 

tests/app/examples/test_cluster_counts.py

@@ -6,6 +6,7 @@
 import importlib.util
 import sys
 from pathlib import Path
+from typing import TypedDict, Unpack
 from unittest.mock import patch
 
 import pytest
@@ -26,8 +27,25 @@
 from firecrown.modeling_tools import ModelingTools
 
 
-def make_builder(tmp_path: Path, **kwargs) -> ExampleClusterCounts:
-    """Create a builder without running the generation workflow."""
+class _BuilderKwargs(TypedDict, total=False):
+    """Optional constructor arguments exercised by these tests."""
+
+    prefix: str
+    dataset: ClusterDataset
+    use_cluster_counts: bool
+    use_mean_log_mass: bool
+    cosmology_spec: Path
+
+
+def make_builder(
+    tmp_path: Path, **kwargs: Unpack[_BuilderKwargs]
+) -> ExampleClusterCounts:
+    """Create a builder without running the generation workflow.
+
+    :param tmp_path: Output directory for the builder.
+    :param kwargs: Optional builder settings used by the test.
+    :returns: Builder with generation suppressed.
+    """
     with patch.object(ExampleClusterCounts, "_proceed_generation"):
         return ExampleClusterCounts(
             output_path=tmp_path, target_framework=Frameworks.COSMOSIS, **kwargs

tests/app/examples/test_cluster_shear.py

@@ -4,6 +4,7 @@
 """
 
 from pathlib import Path
+from typing import TypedDict, Unpack
 from unittest.mock import patch
 
 import numpy as np
@@ -22,8 +23,23 @@
 from firecrown.modeling_tools import ModelingTools
 
 
-def make_builder(tmp_path: Path, **kwargs) -> ExampleClusterShear:
-    """Create a builder without running the generation workflow."""
+class _BuilderKwargs(TypedDict, total=False):
+    """Optional constructor arguments exercised by these tests."""
+
+    prefix: str
+    profile: ClusterProfile
+    cosmology_spec: Path
+
+
+def make_builder(
+    tmp_path: Path, **kwargs: Unpack[_BuilderKwargs]
+) -> ExampleClusterShear:
+    """Create a builder without running the generation workflow.
+
+    :param tmp_path: Output directory for the builder.
+    :param kwargs: Optional builder settings used by the test.
+    :returns: Builder with generation suppressed.
+    """
     with patch.object(ExampleClusterShear, "_proceed_generation"):
         return ExampleClusterShear(
             output_path=tmp_path, target_framework=Frameworks.COSMOSIS, **kwargs

tests/app/examples/test_relocation.py

@@ -30,7 +30,9 @@
     params=CASES,
     ids=[f"{name}-{framework.value}" for name, framework in CASES],
 )
-def fixture_generated(request, tmp_path: Path) -> tuple[Path, Frameworks, str]:
+def fixture_generated(
+    request: pytest.FixtureRequest, tmp_path: Path
+) -> tuple[Path, Frameworks, str]:
     """Generate each example for each framework with the default options."""
     name, framework = request.param
     example_cls, options = EXAMPLES[name]

tests/app/test_experiment.py

@@ -140,7 +140,9 @@ def test_load_experiment_invalid_yaml_raises_exception(
         with pytest.raises(Exception):
             Load(experiment_file=exp_file)
 
-    def test_load_experiment_prints_file_path(self, tmp_path: Path, capsys) -> None:
+    def test_load_experiment_prints_file_path(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that load prints the file path to console."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -201,7 +203,9 @@ def test_view_calls_print_factories(self, tmp_path: Path) -> None:
 class TestViewPrintFactories:
     """Tests for View._print_factories method."""
 
-    def test_print_factories_displays_data_source(self, tmp_path: Path, capsys) -> None:
+    def test_print_factories_displays_data_source(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that _print_factories displays data source information."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -230,7 +234,9 @@ def test_print_factories_displays_data_source(self, tmp_path: Path, capsys) -> N
             assert "Data Source" in captured.out
             assert "data.sacc" in captured.out
 
-    def test_print_factories_displays_filters(self, tmp_path: Path, capsys) -> None:
+    def test_print_factories_displays_filters(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that filters are displayed when present."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -263,7 +269,7 @@ def test_print_factories_displays_filters(self, tmp_path: Path, capsys) -> None:
             assert "scale_cuts.yaml" in captured.out
 
     def test_print_factories_displays_weak_lensing(
-        self, tmp_path: Path, capsys
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
     ) -> None:
         """Test that weak lensing factories are displayed."""
         exp_file = tmp_path / "exp.yaml"
@@ -300,7 +306,7 @@ def test_print_factories_displays_weak_lensing(
             assert "src0" in captured.out
 
     def test_print_factories_displays_number_counts(
-        self, tmp_path: Path, capsys
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
     ) -> None:
         """Test that number counts factories are displayed."""
         exp_file = tmp_path / "exp.yaml"
@@ -335,7 +341,9 @@ def test_print_factories_displays_number_counts(
             assert "NumberCounts" in captured.out
             assert "lens0" in captured.out
 
-    def test_print_factories_displays_cmb(self, tmp_path: Path, capsys) -> None:
+    def test_print_factories_displays_cmb(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that CMB factories are displayed."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -367,7 +375,9 @@ def test_print_factories_displays_cmb(self, tmp_path: Path, capsys) -> None:
             assert "CMBConvergence" in captured.out
             assert "cmb_convergence" in captured.out
 
-    def test_print_factories_displays_systematics(self, tmp_path: Path, capsys) -> None:
+    def test_print_factories_displays_systematics(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that systematics are displayed."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -404,7 +414,7 @@ def test_print_factories_displays_systematics(self, tmp_path: Path, capsys) -> N
             assert "multiplicative_bias" in captured.out
 
     def test_print_factories_empty_systematics_shows_dash(
-        self, tmp_path: Path, capsys
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
     ) -> None:
         """Test that empty systematics list displays a dash."""
         exp_file = tmp_path / "exp.yaml"
@@ -439,7 +449,9 @@ def test_print_factories_empty_systematics_shows_dash(
             # The table should have dash for empty systematics
             assert "WeakLensing" in captured.out
 
-    def test_print_factories_multiple_factories(self, tmp_path: Path, capsys) -> None:
+    def test_print_factories_multiple_factories(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test displaying multiple factories of different types."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {
@@ -491,7 +503,9 @@ def test_print_factories_multiple_factories(self, tmp_path: Path, capsys) -> Non
 class TestFormatSysFunction:
     """Tests for the fmt_sys helper function within _print_factories."""
 
-    def test_fmt_sys_with_multiple_systematics(self, tmp_path: Path, capsys) -> None:
+    def test_fmt_sys_with_multiple_systematics(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture[str]
+    ) -> None:
         """Test that multiple systematics are formatted correctly."""
         exp_file = tmp_path / "exp.yaml"
         exp_data = {