use mocker instead of unittest mock

Author: patricktnast

tests/test_cli.py

@@ -1,10 +1,7 @@
-import json
 from pathlib import Path
 from typing import List
-from unittest.mock import patch
 
 import pytest
-import yaml
 from click.testing import CliRunner
 
 from vivarium_profiling.tools.cli import profile_sim
@@ -47,29 +44,19 @@ def run_profile_sim_command(
     assert result.exit_code == 0
 
 
-def test_profile_sim_parameters():
-    """Test that profile_sim has the expected parameters."""
-    expected_parameters = {
-        "model_specification",
-        "results_directory",
-        "skip_writing",
-        "skip_processing",
-        "profiler",
-    }
-    actual_parameters = {param.name for param in profile_sim.params}
-    assert actual_parameters == expected_parameters
-
-
-@patch("subprocess.run")
-def test_profile_sim_scalene_default(
-    mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
-):
+def test_profile_sim_scalene(mocker, runner: CliRunner, model_spec: str, results_dir: Path):
     """Test profile_sim with scalene profiler (default)."""
     # Mock successful subprocess run
+    mock_subprocess_run = mocker.patch("subprocess.run")
     mock_subprocess_run.return_value = None
 
     # Run the command
-    run_profile_sim_command(runner, model_spec, results_dir)
+    run_profile_sim_command(
+        runner,
+        model_spec,
+        results_dir,
+        extra_args=["--profiler", "scalene"],
+    )
 
     # Verify subprocess was called with correct scalene command
     mock_subprocess_run.assert_called_once()
@@ -89,45 +76,39 @@ def test_profile_sim_scalene_default(
     assert outfile_path.endswith(".json")
 
 
-@patch("subprocess.run")
-@patch("pstats.Stats")
+def test_profile_sim_cprofile_skip_processing(
+    mocker, runner: CliRunner, model_spec: str, results_dir: Path
+):
+    """Test profile_sim with cprofile and processing skipped."""
+    mocker.patch("subprocess.run")
+    run_profile_sim_command(runner, model_spec, results_dir, ["--skip_processing"])
+
+
 def test_profile_sim_cprofile_with_processing(
-    mock_pstats, mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
+    mocker, runner: CliRunner, model_spec: str, results_dir: Path
 ):
     """Test profile_sim with cprofile and post-processing enabled."""
+    mock_subprocess_run = mocker.patch("subprocess.run")
     mock_subprocess_run.return_value = None
 
+    mock_pstats = mocker.patch("pstats.Stats")
     mock_stats_instance = mock_pstats.return_value
 
-    run_profile_sim_command(runner, model_spec, results_dir, ["--profiler", "cprofile"])
+    run_profile_sim_command(runner, model_spec, results_dir)
 
     # Verify that pstats was used for processing
     mock_pstats.assert_called_once()
     mock_stats_instance.sort_stats.assert_called_once_with("cumulative")
     mock_stats_instance.print_stats.assert_called_once()
 
 
-@patch("subprocess.run")
-def test_profile_sim_cprofile_skip_processing(
-    mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
-):
-    """Test profile_sim with cprofile and processing skipped."""
-    mock_subprocess_run.return_value = None
-
-    run_profile_sim_command(
-        runner, model_spec, results_dir, ["--profiler", "cprofile", "--skip_processing"]
-    )
-
-
-@patch("subprocess.run")
 def test_profile_sim_skip_writing(
-    mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
+    mocker, runner: CliRunner, model_spec: str, results_dir: Path
 ):
     """Test profile_sim with skip_writing flag."""
-    # Mock successful subprocess run
-    mock_subprocess_run.return_value = None
+    mock_subprocess_run = mocker.patch("subprocess.run")
+    mocker.patch("pstats.Stats")
 
-    # Run the command with skip_writing
     run_profile_sim_command(runner, model_spec, results_dir, ["--skip_writing"])
 
     # Verify the config override was empty (indicating no output directory set)
@@ -137,12 +118,13 @@ def test_profile_sim_skip_writing(
     assert config_override == "{}"
 
 
-@patch("subprocess.run")
 def test_profile_sim_results_directory_structure(
-    mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
+    mocker, runner: CliRunner, model_spec: str, results_dir: Path
 ):
     """Test that profile_sim creates the expected directory structure and file paths."""
-    mock_subprocess_run.return_value = None
+    mocker.patch("subprocess.run")
+    mocker.patch("pstats.Stats")
+
     run_profile_sim_command(runner, model_spec, results_dir)
 
     # Check that results directory exists and has timestamped subdirectory
@@ -157,13 +139,15 @@ def test_profile_sim_results_directory_structure(
     assert all(part.isdigit() for part in parts)
 
 
-@patch("subprocess.run")
 def test_profile_sim_extra_args(
-    mock_subprocess_run, runner: CliRunner, model_spec: str, results_dir: Path
+    mocker, runner: CliRunner, model_spec: str, results_dir: Path
 ):
     """Test that profile_sim passes extra arguments to scalene."""
-    mock_subprocess_run.return_value = None
-    run_profile_sim_command(runner, model_spec, results_dir, ["--cpu-only", "--html"])
+    mock_subprocess_run = mocker.patch("subprocess.run")
+
+    run_profile_sim_command(
+        runner, model_spec, results_dir, ["--profiler", "scalene", "--cpu-only", "--html"]
+    )
 
     # Verify extra arguments were passed to scalene
     call_args = mock_subprocess_run.call_args[0][0]

tests/test_run_profile.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -15,16 +14,17 @@ def model_spec(test_model_specs: list[Path]) -> str:
     return str(test_model_specs[0])
 
 
-@patch("vivarium_profiling.tools.run_profile.SimulationContext")
-@patch("scalene.scalene_profiler.enable_profiling")
-def test_run_profile_scalene_uses_correct_context_manager(
-    mock_enable_profiling, mock_simulation_context, model_spec: str
-):
+def test_run_profile_scalene_uses_correct_context_manager(mocker, model_spec: str):
     """Test that run_profile_scalene uses scalene's enable_profiling context manager."""
     # Setup mocks
-    mock_sim = MagicMock()
+    mock_sim = mocker.MagicMock()
+    mock_simulation_context = mocker.patch(
+        "vivarium_profiling.tools.run_profile.SimulationContext"
+    )
     mock_simulation_context.return_value = mock_sim
-    mock_context_manager = MagicMock()
+
+    mock_context_manager = mocker.MagicMock()
+    mock_enable_profiling = mocker.patch("scalene.scalene_profiler.enable_profiling")
     mock_enable_profiling.return_value = mock_context_manager
 
     # Run the function
@@ -45,16 +45,19 @@ def test_run_profile_scalene_uses_correct_context_manager(
     mock_sim.run_simulation.assert_called_once()
 
 
-@patch("vivarium_profiling.tools.run_profile.SimulationContext")
-@patch("cProfile.Profile")
 def test_run_profile_cprofile_uses_correct_context_manager(
-    mock_profile_class, mock_simulation_context, model_spec: str, tmp_path: Path
+    mocker, model_spec: str, tmp_path: Path
 ):
     """Test that run_profile_cprofile uses cProfile.Profile context manager."""
     # Setup mocks
-    mock_sim = MagicMock()
+    mock_sim = mocker.MagicMock()
+    mock_simulation_context = mocker.patch(
+        "vivarium_profiling.tools.run_profile.SimulationContext"
+    )
     mock_simulation_context.return_value = mock_sim
-    mock_profiler = MagicMock()
+
+    mock_profiler = mocker.MagicMock()
+    mock_profile_class = mocker.patch("cProfile.Profile")
     mock_profile_class.return_value = mock_profiler
 
     # The context manager should return the same profiler instance