Merge pull request #453 from The-Strategy-Unit/fix_config

Author: tomjemmett

src/nhp/docker/__main__.py

@@ -5,7 +5,7 @@
 import os
 import threading
 
-from nhp.docker import config
+from nhp.docker.config import Config
 from nhp.docker.run import RunWithAzureStorage, RunWithLocalStorage
 from nhp.model.run import run_all
 
@@ -37,7 +37,7 @@ def parse_args():
     return parser.parse_args()
 
 
-def main():
+def main(config: Config = Config()):
     """The main method."""
     args = parse_args()
 
@@ -50,7 +50,7 @@ def main():
     if args.local_storage:
         runner = RunWithLocalStorage(args.params_file)
     else:
-        runner = RunWithAzureStorage(args.params_file, config.APP_VERSION)
+        runner = RunWithAzureStorage(args.params_file, config)
 
     logging.info("running model for: %s", args.params_file)
     logging.info("container timeout: %ds", config.CONTAINER_TIMEOUT_SECONDS)
@@ -73,12 +73,13 @@ def init():
     """Method for calling main."""
     if __name__ == "__main__":
         exc = None
+        config = Config()
         try:
             # start a timer to kill the container if we reach a timeout
             t = threading.Timer(config.CONTAINER_TIMEOUT_SECONDS, _exit_container)
             t.start()
             # run the model
-            main()
+            main(config)
         except Exception as e:
             logging.error("An error occurred: %s", str(e))
             exc = e

src/nhp/docker/config.py

@@ -1,16 +1,44 @@
 """config values for docker container."""
 
+import os
+
 import dotenv
 
-__config_values = dotenv.dotenv_values()
 
+class Config:
+    """Configuration class for Docker container."""
+
+    __DEFAULT_CONTAINER_TIMEOUT_SECONDS = 60 * 60  # 1 hour
+
+    def __init__(self):
+        """Configuration settings for the Docker container."""
+        dotenv.load_dotenv()
+
+        self._app_version = os.environ.get("APP_VERSION", "dev")
+        self._data_version = os.environ.get("DATA_VERSION", "dev")
+        self._storage_account = os.environ.get("STORAGE_ACCOUNT")
+
+        self._container_timeout_seconds = os.environ.get("CONTAINER_TIMEOUT_SECONDS")
+
+    @property
+    def APP_VERSION(self) -> str:
+        """What is the version of the app?"""
+        return self._app_version
 
-APP_VERSION: str = __config_values.get("APP_VERSION", "dev")  # type: ignore
-DATA_VERSION: str = __config_values.get("DATA_VERSION", "dev")  # type: ignore
+    @property
+    def DATA_VERSION(self) -> str:
+        """What version of the data are we using?"""
+        return self._data_version
 
-STORAGE_ACCOUNT: str | None = __config_values.get("STORAGE_ACCOUNT", None)
+    @property
+    def STORAGE_ACCOUNT(self) -> str:
+        """What is the name of the storage account?"""
+        if self._storage_account is None:
+            raise ValueError("STORAGE_ACCOUNT environment variable must be set")
+        return self._storage_account
 
-__DEFAULT_CONTAINER_TIMEOUT_SECONDS = 60 * 60  # 1 hour
-CONTAINER_TIMEOUT_SECONDS = int(
-    __config_values.get("CONTAINER_TIMEOUT_SECONDS", __DEFAULT_CONTAINER_TIMEOUT_SECONDS)  # type: ignore
-)
+    @property
+    def CONTAINER_TIMEOUT_SECONDS(self) -> int:
+        """How long should the container run before timing out?"""
+        t = self._container_timeout_seconds
+        return self.__DEFAULT_CONTAINER_TIMEOUT_SECONDS if t is None else int(t)

src/nhp/docker/run.py

@@ -12,7 +12,7 @@
 from azure.storage.blob import BlobServiceClient
 from azure.storage.filedatalake import DataLakeServiceClient
 
-from nhp.docker import config
+from nhp.docker.config import Config
 from nhp.model.helpers import load_params
 from nhp.model.run import noop_progress_callback
 
@@ -51,27 +51,35 @@ def progress_callback(self) -> Callable[[Any], Callable[[Any], None]]:
 class RunWithAzureStorage:
     """Methods for running with azure storage."""
 
-    def __init__(self, filename: str, app_version: str = "dev"):
+    def __init__(self, filename: str, config: Config = Config()):
         """Initialise RunWithAzureStorage.
 
         :param filename:
         :type filename: str
-        :param app_version: the version of the app, where we will load data from. defaults to "dev"
-        :type app_version: str, optional
+        :param config: The configuration for the run
+        :type config: Config
         """
         logging.getLogger("azure.storage.common.storageclient").setLevel(logging.WARNING)
         logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
             logging.WARNING
         )
+        self._config = config
 
-        self._app_version = re.sub("(\\d+\\.\\d+)\\..*", "\\1", app_version)
+        self._app_version = re.sub("(\\d+\\.\\d+)\\..*", "\\1", config.APP_VERSION)
+
+        self._blob_storage_account_url = (
+            f"https://{self._config.STORAGE_ACCOUNT}.blob.core.windows.net"
+        )
+        self._adls_storage_account_url = (
+            f"https://{self._config.STORAGE_ACCOUNT}.dfs.core.windows.net"
+        )
 
         self.params = self._get_params(filename)
         self._get_data(self.params["start_year"], self.params["dataset"])
 
     def _get_container(self, container_name: str):
         return BlobServiceClient(
-            account_url=f"https://{config.STORAGE_ACCOUNT}.blob.core.windows.net",
+            account_url=self._blob_storage_account_url,
             credential=DefaultAzureCredential(),
         ).get_container_client(container_name)
 
@@ -103,11 +111,11 @@ def _get_data(self, year: str, dataset: str) -> None:
         """
         logging.info("downloading data (%s / %s)", year, dataset)
         fs_client = DataLakeServiceClient(
-            account_url=f"https://{config.STORAGE_ACCOUNT}.dfs.core.windows.net",
+            account_url=self._adls_storage_account_url,
             credential=DefaultAzureCredential(),
         ).get_file_system_client("data")
 
-        version = config.DATA_VERSION
+        version = self._config.DATA_VERSION
 
         paths = [p.name for p in fs_client.get_paths(version, recursive=False)]
 

tests/unit/nhp/docker/test___main__.py

@@ -1,11 +1,10 @@
 """test docker run."""
 
 import time
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 
-from nhp.docker import config
 from nhp.docker.__main__ import main, parse_args
 
 
@@ -93,6 +92,12 @@ def test_main_azure(mocker):
     rwls = mocker.patch("nhp.docker.__main__.RunWithLocalStorage")
     rwas = mocker.patch("nhp.docker.__main__.RunWithAzureStorage")
 
+    config = Mock()
+    config.APP_VERSION = "dev"
+    config.DATA_VERSION = "dev"
+    config.CONTAINER_TIMEOUT_SECONDS = 3600
+    config.STORAGE_ACCOUNT = "sa"
+
     params = {
         "model_runs": 256,
         "start_year": 2019,
@@ -108,11 +113,11 @@ def test_main_azure(mocker):
     )
 
     # act
-    main()
+    main(config)
 
     # assert
     rwls.assert_not_called()
-    rwas.assert_called_once_with("params.json", "dev")
+    rwas.assert_called_once_with("params.json", config)
 
     s = rwas()
     ru_m.assert_called_once_with(params, "data", s.progress_callback(), False)
@@ -121,6 +126,9 @@ def test_main_azure(mocker):
 
 def test_init(mocker):
     """It should run the main method if __name__ is __main__."""
+    config = mocker.patch("nhp.docker.__main__.Config")
+    config().CONTAINER_TIMEOUT_SECONDS = 3600
+
     import nhp.docker.__main__ as r
 
     main_mock = mocker.patch("nhp.docker.__main__.main")
@@ -130,31 +138,33 @@ def test_init(mocker):
 
     with patch.object(r, "__name__", "__main__"):
         r.init()  # should call main
-        main_mock.assert_called_once()
+        main_mock.assert_called_once_with(config())
 
 
 def test_init_timeout_call_exit(mocker):
-    config.CONTAINER_TIMEOUT_SECONDS = 0.1
+    config = mocker.patch("nhp.docker.__main__.Config")
+    config().CONTAINER_TIMEOUT_SECONDS = 0.1
 
     import nhp.docker.__main__ as r
 
     main_mock = mocker.patch("nhp.docker.__main__.main")
     exit_container_mock = mocker.patch("nhp.docker.__main__._exit_container")
-    main_mock.side_effect = lambda: time.sleep(0.2)
+    main_mock.side_effect = lambda *args, **kwargs: time.sleep(0.2)
     with patch.object(r, "__name__", "__main__"):
         r.init()
 
     exit_container_mock.assert_called_once()
 
 
 def test_init_timeout_dont_call_exit(mocker):
-    config.CONTAINER_TIMEOUT_SECONDS = 0.1
-
     import nhp.docker.__main__ as r
 
+    config = mocker.patch("nhp.docker.__main__.Config")
+    config().CONTAINER_TIMEOUT_SECONDS = 0.1
+
     main_mock = mocker.patch("nhp.docker.__main__.main")
     exit_container_mock = mocker.patch("nhp.docker.__main__._exit_container")
-    main_mock.side_effect = lambda: time.sleep(0.02)
+    main_mock.side_effect = lambda *args, **kwargs: time.sleep(0.02)
     with patch.object(r, "__name__", "__main__"):
         r.init()
 

tests/unit/nhp/docker/test_config.py

@@ -1,51 +1,57 @@
-import importlib
+import os
 from unittest.mock import patch
 
-import nhp.docker.config as c
+import pytest
+
+from nhp.docker.config import Config
 
 
 def test_config_sets_values_from_envvars(mocker):
     # arrange
-    mocker.patch(
-        "dotenv.dotenv_values",
-        return_value={
+    mocker.patch("dotenv.load_dotenv")
+
+    # act
+    with patch.dict(
+        os.environ,
+        {
             "APP_VERSION": "app version",
             "DATA_VERSION": "data version",
             "STORAGE_ACCOUNT": "storage account",
             "CONTAINER_TIMEOUT_SECONDS": "123",
         },
-    )
-
-    # act
-    importlib.reload(c)
+    ):
+        config = Config()
 
     # assert
-    assert c.APP_VERSION == "app version"
-    assert c.DATA_VERSION == "data version"
-    assert c.STORAGE_ACCOUNT == "storage account"
-    assert c.CONTAINER_TIMEOUT_SECONDS == 123
+    assert config.APP_VERSION == "app version"
+    assert config.DATA_VERSION == "data version"
+    assert config.STORAGE_ACCOUNT == "storage account"
+    assert config.CONTAINER_TIMEOUT_SECONDS == 123
 
 
 def test_config_uses_default_values(mocker):
     # arrange
-    mocker.patch("dotenv.dotenv_values", return_value={})
+    mocker.patch("dotenv.load_dotenv")
 
     # act
-    importlib.reload(c)
+    config = Config()
 
     # assert
-    assert c.APP_VERSION == "dev"
-    assert c.DATA_VERSION == "dev"
-    assert not c.STORAGE_ACCOUNT
-    assert c.CONTAINER_TIMEOUT_SECONDS == 3600
+    assert config.APP_VERSION == "dev"
+    assert config.DATA_VERSION == "dev"
+
+    with pytest.raises(ValueError, match="STORAGE_ACCOUNT environment variable must be set"):
+        config.STORAGE_ACCOUNT
+
+    assert config.CONTAINER_TIMEOUT_SECONDS == 3600
 
 
 def test_config_calls_dotenv_load(mocker):
     # arrange
-    m = mocker.patch("dotenv.dotenv_values", return_value={})
+    m = mocker.patch("dotenv.load_dotenv")
 
     # act
-    importlib.reload(c)
+    config = Config()
 
     # assert
     m.assert_called_once()

tests/unit/nhp/docker/test_run.py

@@ -59,14 +59,25 @@ def mock_run_with_azure_storage():
     rwas.params = {}
     rwas._app_version = "dev"
 
+    rwas._config = Mock()
+    rwas._config.APP_VERSION = "dev"
+    rwas._config.DATA_VERSION = "dev"
+    rwas._config.STORAGE_ACCOUNT = "sa"
+    rwas._config.CONTAINER_TIMEOUT_SECONDS = 3600
+
+    rwas._blob_storage_account_url = "https://sa.blob.core.windows.net"
+    rwas._adls_storage_account_url = "https://sa.dfs.core.windows.net"
+
     return rwas
 
 
-@pytest.mark.parametrize(
-    "args, expected_version", [(["filename"], "dev"), (["filename", "v0.3.5"], "v0.3")]
-)
-def test_RunWithAzureStorage_init(mocker, args, expected_version):
+@pytest.mark.parametrize("actual_version, expected_version", [("dev", "dev"), ("v0.3.5", "v0.3")])
+def test_RunWithAzureStorage_init(mocker, actual_version, expected_version):
     # arrange
+    config = mocker.patch("nhp.docker.run.Config")
+    config().APP_VERSION = actual_version
+    config().STORAGE_ACCOUNT = "sa"
+
     expected_params = {"start_year": 2020, "dataset": "synthetic"}
     gpm = mocker.patch(
         "nhp.docker.run.RunWithAzureStorage._get_params",
@@ -75,11 +86,13 @@ def test_RunWithAzureStorage_init(mocker, args, expected_version):
     gdm = mocker.patch("nhp.docker.run.RunWithAzureStorage._get_data")
 
     # act
-    s = RunWithAzureStorage(*args)  # type: ignore
+    s = RunWithAzureStorage("filename", config())
 
     # assert
     assert s._app_version == expected_version
     assert s.params == expected_params
+    assert s._blob_storage_account_url == "https://sa.blob.core.windows.net"
+    assert s._adls_storage_account_url == "https://sa.dfs.core.windows.net"
 
     gpm.assert_called_once_with("filename")
 
@@ -97,7 +110,6 @@ def test_RunWithAzureStorage_get_container(mock_run_with_azure_storage, mocker):
 
     dac_m = mocker.patch("nhp.docker.run.DefaultAzureCredential", return_value="cred")
 
-    config.STORAGE_ACCOUNT = "sa"
     # act
     actual = s._get_container("container")
 
@@ -163,8 +175,6 @@ def paths_helper(i):
 
     dac_m = mocker.patch("nhp.docker.run.DefaultAzureCredential", return_value="cred")
 
-    config.STORAGE_ACCOUNT = "sa"
-
     # act
     with patch("builtins.open", mock_open()) as mock_file:
         s._get_data(2020, "synthetic")