test(simulation): add more unit and functional tests based on the template

Author: amyheather

tests/test_functionaltest.py

@@ -5,11 +5,15 @@
 functionality.
 """
 
+from joblib import cpu_count
+import numpy as np
+import pandas as pd
 import pytest
 
 from simulation.parameters import (ASUArrivals, RehabArrivals, ASULOS,
                                    RehabLOS, Param)
 from simulation.model import Model
+from simulation.runner import Runner
 
 
 @pytest.mark.parametrize("warm_up_period", [(0), (1)])
@@ -117,3 +121,238 @@ def test_long_los(stroke_no_esd_mean):
 
 
 # TODO: test_ev_5/6/7/8 - requires patient counts by unit and patient type
+
+def test_warmup_only():
+    """
+    Check that results are as expected if model is run with no data collection
+    period, and only a warm-up period.
+
+    Notes
+    -----
+    Inspired by `test_warmup_only` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    # Run with warm-up period only
+    param = Param(warm_up_period=500, data_collection_period=0)
+    model = Model(param=param, run_number=0)
+    model.run()
+
+    # Check that there are no entries for patients or audit list
+    assert len(model.patients) == 0
+    assert len(model.audit_list) == 0
+
+    # But that current occupancy is positive, as it reflects patients who
+    # arrived during warm-up still occupying space
+    assert model.asu_occupancy > 0
+    assert model.rehab_occupancy > 0
+
+
+def test_warmup_impact():
+    """
+    Check that running with warm-up leads to different results than without.
+
+    Notes
+    -----
+    Adapted from `test_warmup_impact` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    def helper_warmup(warm_up_period):
+        """
+        Helper function to run model with high arrivals and specified warm-up.
+
+        Arguments:
+            warm_up_period (float):
+                Duration of the warm-up period - running simulation but not yet
+                collecting results.
+        """
+        param = Param(warm_up_period=warm_up_period,
+                      data_collection_period=1500)
+        model = Model(param=param, run_number=0)
+        model.run()
+        return model
+
+    # Run model with and without warm-up period
+    model_warmup = helper_warmup(warm_up_period=500)
+    model_none = helper_warmup(warm_up_period=0)
+
+    # Extract result of first patient
+    first_warmup = model_warmup.patients[0]
+    first_none = model_none.patients[0]
+
+    # Check that model with warm-up has arrival time later than warm-up length
+    assert first_warmup.asu_arrival_time > 500, (
+        "Expect first patient to arrive after time 500 when model is run " +
+        f"with warm-up (length 500), but got {first_warmup.asu_arrival_time}."
+    )
+
+    # Check that model without warm-up has arrival first arrival after time 0
+    assert first_none.asu_arrival_time > 0, (
+        "Expect first patient to arrive after time 0 when model is run " +
+        f"without warm-up, but got {first_none.asu_arrival_time}."
+    )
+
+    # Check that the first interval audit entry with no warm-up has occupancy
+    # 0, but first entry with warm-up has occupancy > 0
+    assert model_warmup.audit_list[0]["asu_occupancy"] > 0
+    assert model_warmup.audit_list[0]["rehab_occupancy"] > 0
+    assert model_none.audit_list[0]["asu_occupancy"] == 0
+    assert model_none.audit_list[0]["rehab_occupancy"] == 0
+
+    # Check time of first entry in the audit lists
+    assert model_warmup.audit_list[0]["time"] == 500
+    assert model_none.audit_list[0]["time"] == 0
+
+
+def test_changing_occupancy():
+    """
+    Test that adjusting parameters alters the observed occupancy.
+
+    Notes
+    -----
+    Inspired by `test_arrivals_decrease` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    # Run model with lots of stroke arrivals (IAT 1) and fewer (IAT 100)
+    initial_model = Model(Param(ASUArrivals(stroke=1)), run_number=0)
+    initial_model.run()
+    adj_model = Model(Param(ASUArrivals(stroke=100)), run_number=0)
+    adj_model.run()
+
+    # Check that patient list is longer with more arrivals
+    i_n = len(initial_model.patients)
+    a_n = len(adj_model.patients)
+    assert i_n > a_n, (
+        "Expect high IAT to have more patients than low IAT, but " +
+        f"saw {i_n} for high and {a_n} for low"
+    )
+
+    # Check that sum of the audit occupancies is higher with more arrivals
+    initial_audit = pd.DataFrame(initial_model.audit_list)
+    adj_audit = pd.DataFrame(adj_model.audit_list)
+
+    i_asu = initial_audit["asu_occupancy"].sum()
+    a_asu = adj_audit["asu_occupancy"].sum()
+    assert i_asu > a_asu, (
+        "Expect high IAT to have higher ASU occupancy than low IAT, but " +
+        f"saw {i_asu} for high and {a_asu} for low")
+
+    i_reh = initial_audit["rehab_occupancy"].sum()
+    a_reh = adj_audit["rehab_occupancy"].sum()
+    assert i_reh > a_reh, (
+        "Expect high IAT to have higher rehab occupancy than low IAT, but " +
+        f"saw {i_reh} for high and {a_reh} for low")
+
+
+def test_seed_stability():
+    """
+    Check that two runs using the same random seed return the same results.
+
+    Notes
+    -----
+    Adapted from `seed_seed_stability` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    # Run model twice, with same run number (and therefore same seed) each time
+    runner1 = Runner(param=Param())
+    result1 = runner1.run_single(run=33)
+    runner2 = Runner(param=Param())
+    result2 = runner2.run_single(run=33)
+
+    # Check that the dataframes are equal
+    pd.testing.assert_frame_equal(result1["asu"], result2["asu"])
+    pd.testing.assert_frame_equal(result1["rehab"], result2["rehab"])
+
+
+def test_parallel():
+    """
+    Check that sequential and parallel execution produce consistent results.
+
+    Notes
+    -----
+    Adapted from `test_parallel` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    # Sequential (1 core) and parallel (-1 cores) execution
+    results = {}
+    for mode, cores in [("seq", 1), ("par", -1)]:
+        param = Param(cores=cores)
+        runner = Runner(param)
+        results[mode] = runner.run_single(run=0)
+
+    # Verify results are identical
+    pd.testing.assert_frame_equal(results["seq"]["asu"],
+                                  results["par"]["asu"])
+    pd.testing.assert_frame_equal(results["seq"]["rehab"],
+                                  results["par"]["rehab"])
+
+
+@pytest.mark.parametrize("cores", [
+    (-2), (0), (cpu_count()), (cpu_count()+1)
+])
+def test_valid_cores(cores):
+    """
+    Check there is error handling for input of invalid number of cores.
+
+    Notes
+    -----
+    Copied from github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    param = Param(cores=cores)
+    runner = Runner(param)
+    with pytest.raises(ValueError):
+        runner.run_reps()
+
+
+def test_sampled_distributions():
+    """
+    Ensure that the mean of sampled values from arrival_dist, los_dist, and
+    routing_dist are close to their expected values.
+
+    Notes
+    -----
+    Adapted from `test_sampled_times` in
+    github.com/pythonhealthdatascience/rap_template_python_des/.
+    """
+    param = Param()
+    model = Model(param, run_number=0)
+
+    # Test arrival_dist
+    for unit in ['asu', 'rehab']:
+        for patient_type, dist in model.arrival_dist[unit].items():
+            samples = [dist.sample() for _ in range(10000)]
+            observed_mean = np.mean(samples)
+            expected_mean = getattr(
+                getattr(param, f"{unit}_arrivals"), patient_type)
+            assert np.isclose(observed_mean, expected_mean, rtol=0.05), (
+                f"Expected mean arrival time for {unit} {patient_type} ≈ "
+                f"{expected_mean}, but got {observed_mean}.")
+
+    # Test los_dist
+    for unit in ['asu', 'rehab']:
+        for patient_type, dist in model.los_dist[unit].items():
+            samples = [dist.sample() for _ in range(10000)]
+            observed_mean = np.mean(samples)
+            expected_mean = getattr(
+                getattr(param, f"{unit}_los"), patient_type)["mean"]
+            assert np.isclose(observed_mean, expected_mean, rtol=0.05), (
+                f"Expected mean LOS for {unit} {patient_type} ≈ "
+                f"{expected_mean}, but got {observed_mean}.")
+
+    # Test routing_dist
+    for unit in ['asu', 'rehab']:
+        for patient_type, dist in model.routing_dist[unit].items():
+            samples = [dist.sample() for _ in range(10000)]
+            observed_probs = {dest: samples.count(dest) / len(samples)
+                              for dest in set(samples)}
+            expected_probs = getattr(
+                getattr(param, f"{unit}_routing"), patient_type)
+            for dest in expected_probs:
+                # If a destination has a very low probability, it might not
+                # appear in samples. In that case, set the observed probability
+                # to 0
+                observed_prob = observed_probs.get(dest, 0)
+                assert np.isclose(observed_prob,
+                                  expected_probs[dest], atol=0.05), (
+                    f"Expected routing probability for {unit} {patient_type} "
+                    f"to {dest} ≈ {expected_probs[dest]}, but got "
+                    f"{observed_prob}.")

tests/test_unittest.py

@@ -6,6 +6,11 @@
 work as intended.
 """
 
+from io import StringIO
+import logging
+import os
+from unittest.mock import patch, MagicMock
+
 import numpy as np
 import pytest
 from sim_tools.distributions import Exponential, Lognormal, Discrete
@@ -15,6 +20,7 @@
     ASURouting, RehabRouting, Param)
 from simulation.model import Model
 from simulation.runner import Runner
+from simulation.logging import SimLogger
 
 
 # -----------------------------------------------------------------------------
@@ -243,3 +249,62 @@ def test_get_occupancy_freq():
 
     # Check prob_delay calculation
     assert np.allclose(result_df["prob_delay"], expected_prob_delay)
+
+
+# -----------------------------------------------------------------------------
+# SimLogger
+# -----------------------------------------------------------------------------
+
+def test_log_to_console():
+    """
+    Confirm that logger.log() prints the provided message to the console.
+
+    Notes
+    -----
+    Test from github.com/pythonhealthdatascience/rap_template_python_des.
+    """
+    with patch('sys.stdout', new_callable=StringIO) as mock_stdout:
+        logger = SimLogger(log_to_console=True)
+        logger.log(sim_time=None, msg="Test console log")
+        # Check if console output matches
+        assert "Test console log" in mock_stdout.getvalue()
+
+
+def test_log_to_file():
+    """
+    Confirm that logger.log() would output the message to a .log file at the
+    provided file path.
+
+    Notes
+    -----
+    Test from github.com/pythonhealthdatascience/rap_template_python_des.
+    """
+    # Mock the file open operation
+    with patch('builtins.open', new_callable=MagicMock) as mock_open:
+        # Create the logger and log a simple example
+        logger = SimLogger(log_to_file=True, file_path="test.log")
+        logger.log(sim_time=None, msg="Log message")
+
+        # Check that the file was opened in write mode at the absolute path
+        mock_open.assert_called_with(
+            os.path.abspath("test.log"), "w", encoding="locale", errors=None)
+
+        # Verify a FileHandler is attached to the logger
+        assert (any(isinstance(handler, logging.FileHandler)
+                    for handler in logger.logger.handlers))
+
+
+def test_invalid_path():
+    """
+    Ensure there is appropriate error handling for an invalid file path.
+    """
+    with pytest.raises(ValueError):
+        SimLogger(log_to_file=True, file_path="/invalid/path/to/log.log")
+
+
+def test_invalid_file_extension():
+    """
+    Ensure there is appropriate error handling for an invalid file extension.
+    """
+    with pytest.raises(ValueError):
+        SimLogger(log_to_file=True, file_path="test.txt")