Merge pull request #482 from The-Strategy-Unit/alter_run_all

allows run_all to use different data sources

Author: tomjemmett

src/nhp/docker/__main__.py

@@ -2,10 +2,10 @@
 
 import argparse
 import logging
-import os
 
 from nhp.docker.config import Config
 from nhp.docker.run import RunWithAzureStorage, RunWithLocalStorage
+from nhp.model.data import Local
 from nhp.model.run import run_all
 
 
@@ -54,7 +54,10 @@ def main(config: Config = Config()):
     logging.info("app_version:  %s", runner.params["app_version"])
 
     saved_files, results_file = run_all(
-        runner.params, "data", runner.progress_callback(), args.save_full_model_results
+        runner.params,
+        Local.create("data"),
+        runner.progress_callback(),
+        args.save_full_model_results,
     )
 
     runner.finish(results_file, saved_files, args.save_full_model_results)

src/nhp/model/__main__.py

@@ -16,6 +16,7 @@
 import logging
 
 from nhp.model.aae import AaEModel
+from nhp.model.data import Local
 from nhp.model.helpers import load_params
 from nhp.model.inpatients import InpatientsModel
 from nhp.model.outpatients import OutpatientsModel
@@ -65,7 +66,7 @@ def main() -> None:
 
             run_all(
                 params,
-                args.data_path,
+                Local.create(args.data_path),
                 lambda _: lambda _: None,
                 args.save_full_model_results,
             )

src/nhp/model/run.py

@@ -111,7 +111,7 @@ def noop_progress_callback(_: Any) -> Callable[[Any], None]:
 
 def run_all(
     params: dict,
-    data_path: str,
+    nhp_data: Callable[[int, str], Data],
     progress_callback: Callable[[Any], Callable[[Any], None]] = noop_progress_callback,
     save_full_model_results: bool = False,
 ) -> Tuple[list, str]:
@@ -121,8 +121,8 @@ def run_all(
 
     :param params: the parameters to use for this model run
     :type params: dict
-    :param data_path: where the data is stored
-    :type data_path: str
+    :param nhp_data: the Data class to use for loading data
+    :type nhp_data: Callable[[int, str], Data]
     :param progress_callback: a callback function for updating progress.
         Defaults to noop_progress_callback.
     :type progress_callback: Callable[[str], Callable[[Any], Any]]
@@ -134,8 +134,6 @@ def run_all(
     model_types = [InpatientsModel, OutpatientsModel, AaEModel]
     run_params = Model.generate_run_params(params)
 
-    nhp_data = Local.create(data_path)
-
     # set the data path in the HealthStatusAdjustment class
     hsa = HealthStatusAdjustmentInterpolated(
         nhp_data(params["start_year"], params["dataset"]), params["start_year"]

tests/integration/nhp/model/test_run_model.py

@@ -61,8 +61,9 @@ def test_all_model_runs(params_path, data_dir):
     params = load_params(params_path)
     params["model_runs"] = 4
 
+    nhp_data = Local.create(data_dir)
     # act
-    actual = run_all(params, data_dir)
+    actual = run_all(params, nhp_data)
 
     # assert
     res_path = "results/synthetic/test/20220101_000000"

tests/unit/nhp/docker/test___main__.py

@@ -46,6 +46,9 @@ def test_main_local(mocker):
     rwls = mocker.patch("nhp.docker.__main__.RunWithLocalStorage")
     rwas = mocker.patch("nhp.docker.__main__.RunWithAzureStorage")
 
+    local_data_mock = mocker.patch("nhp.docker.__main__.Local")
+    local_data_mock.create.return_value = "data"
+
     params = {
         "model_runs": 256,
         "start_year": 2019,
@@ -71,6 +74,8 @@ def test_main_local(mocker):
     ru_m.assert_called_once_with(params, "data", s.progress_callback(), False)
     s.finish.assert_called_once_with("results.json", "list_of_results", False)
 
+    local_data_mock.create.assert_called_once_with("data")
+
 
 def test_main_azure(mocker):
     # arrange
@@ -82,6 +87,9 @@ def test_main_azure(mocker):
     rwls = mocker.patch("nhp.docker.__main__.RunWithLocalStorage")
     rwas = mocker.patch("nhp.docker.__main__.RunWithAzureStorage")
 
+    local_data_mock = mocker.patch("nhp.docker.__main__.Local")
+    local_data_mock.create.return_value = "data"
+
     config = Mock()
     config.APP_VERSION = "dev"
     config.DATA_VERSION = "dev"
@@ -112,6 +120,8 @@ def test_main_azure(mocker):
     ru_m.assert_called_once_with(params, "data", s.progress_callback(), False)
     s.finish.assert_called_once_with("results.json", "list_of_results", False)
 
+    local_data_mock.create.assert_called_once_with("data")
+
 
 def test_init(mocker):
     """It should run the main method if __name__ is __main__."""

tests/unit/nhp/model/test__main__.py

@@ -68,6 +68,8 @@ def test_main_all_runs(mocker):
     args.save_full_model_results = False
     mocker.patch("nhp.model.__main__._parse_args", return_value=args)
     ldp_mock = mocker.patch("nhp.model.__main__.load_params", return_value="params")
+    local_data_mock = mocker.patch("nhp.model.__main__.Local")
+    local_data_mock.create.return_value = "data"
 
     run_all_mock = mocker.patch("nhp.model.__main__.run_all")
     run_single_mock = mocker.patch("nhp.model.__main__.run_single_model_run")
@@ -83,6 +85,7 @@ def test_main_all_runs(mocker):
 
     run_single_mock.assert_not_called()
     ldp_mock.assert_called_once_with("queue/params.json")
+    local_data_mock.create.assert_called_once_with("data")
 
 
 def test_init(mocker):

tests/unit/nhp/model/test_run.py

@@ -87,7 +87,6 @@ def test_run_all(mocker):
     )
     gr_m = mocker.patch("nhp.model.run.generate_results_json", return_value="results_json_path")
     sr_m = mocker.patch("nhp.model.run.save_results_files", return_value="results_paths")
-    nd_m = mocker.patch("nhp.model.run.Local")
 
     pc_m = Mock()
     pc_m().return_value = "progress callback"
@@ -102,16 +101,15 @@ def test_run_all(mocker):
         "model_runs": 10,
         "create_datetime": "20230123_012345",
     }
+    data_mock = Mock(return_value="nhp_data")
 
     # act
-    actual = run_all(params, "data_path", pc_m, False)
+    actual = run_all(params, data_mock, pc_m, False)
 
     # assert
     assert actual == ("results_paths", "results_json_path")
 
-    nd_m.create.assert_called_once_with("data_path")
-    nd_c = nd_m.create()
-    nd_c.assert_called_once_with(2020, "synthetic")
+    data_mock.assert_called_once_with(2020, "synthetic")
 
     assert pc_m.call_args_list == [
         call("Inpatients"),
@@ -120,13 +118,13 @@ def test_run_all(mocker):
     ]
 
     grp_m.assert_called_once_with(params)
-    hsa_m.assert_called_once_with(nd_c(2020, "synthetic"), 2020)
+    hsa_m.assert_called_once_with("nhp_data", 2020)
 
     assert rm_m.call_args_list == [
         call(
             m,
             params,
-            nd_c,
+            data_mock,
             "hsa",
             {"variant": "variants"},
             pc_m(),
@@ -136,15 +134,12 @@ def test_run_all(mocker):
     ]
 
     cr_m.assert_called_once_with(["ip", "op", "aae"])
-    # weaker form of checking, but as we intended to drop this function in the future don't expend
-    # effort to fixing this part of the test
-    gr_m.assert_called_once()
-    # gr_m.assert_called_once_with(
-    #     {"default": "combined_results"},
-    #     "combined_step_counts",
-    #     params,
-    #     {"variant": "variants"},
-    # )
+    gr_m.assert_called_once_with(
+        {"default": "combined_results", "step_counts": "combined_step_counts"},
+        "combined_step_counts",
+        params,
+        {"variant": "variants"},
+    )
     sr_m.assert_called_once_with(
         {"default": "combined_results", "step_counts": "combined_step_counts"}, params
     )