From 6fa419a09d40c7559244ce21916e1f9c96bdf40d Mon Sep 17 00:00:00 2001 From: pciturri Date: Tue, 30 Jul 2024 21:37:47 +0200 Subject: [PATCH] ft: Created environments module, which handle the creation, staging of conda/pip/docker environments, as well as the running a source code within. tests: added tests for EnvironmentHandler class and subclasses build: added packaging as requirement. Harmonized requirements.txt and requirements_dev.txt --- .../case_g/models/pymock_nb/pyproject.toml | 0 floatcsep/environments.py | 469 ++++++++++++++++++ floatcsep/model.py | 34 +- floatcsep/registry.py | 3 +- pyproject.toml | 7 +- requirements.txt | 2 + requirements_dev.txt | 3 +- setup.cfg | 3 + tests/artifacts/models/td_model/setup.cfg | 0 tests/unit/test_environments.py | 365 ++++++++++++++ 10 files changed, 861 insertions(+), 25 deletions(-) create mode 100644 examples/case_g/models/pymock_nb/pyproject.toml create mode 100644 floatcsep/environments.py create mode 100644 tests/artifacts/models/td_model/setup.cfg create mode 100644 tests/unit/test_environments.py diff --git a/examples/case_g/models/pymock_nb/pyproject.toml b/examples/case_g/models/pymock_nb/pyproject.toml new file mode 100644 index 0000000..e69de29 diff --git a/floatcsep/environments.py b/floatcsep/environments.py new file mode 100644 index 0000000..96f3bab --- /dev/null +++ b/floatcsep/environments.py @@ -0,0 +1,469 @@ +import logging +import sys +import os +import subprocess +import hashlib +import shutil +import venv +import configparser +from abc import ABC, abstractmethod +from packaging.specifiers import SpecifierSet + +log = logging.getLogger("floatLogger") + + +class EnvironmentManager(ABC): + """ + Abstract base class for managing different types of environments. + This class defines the interface for creating, checking existence, + running commands, and installing dependencies in various environment types. + """ + + @abstractmethod + def __init__(self, base_name: str, model_directory: str): + """ + Initializes the environment manager with a base name and model directory. + + Args: + base_name (str): The base name for the environment. + model_directory (str): The directory containing the model files. + """ + self.base_name = base_name + self.model_directory = model_directory + + @abstractmethod + def create_environment(self, force=False): + """ + Creates the environment. If 'force' is True, it will remove any existing + environment with the same name before creating a new one. + + Args: + force (bool): Whether to forcefully remove an existing environment. + """ + pass + + @abstractmethod + def env_exists(self): + """ + Checks if the environment already exists. + + Returns: + bool: True if the environment exists, False otherwise. + """ + pass + + @abstractmethod + def run_command(self, command): + """ + Executes a command within the context of the environment. + + Args: + command (str): The command to be executed. + """ + pass + + @abstractmethod + def install_dependencies(self): + """ + Installs the necessary dependencies for the environment based on the + specified configuration or requirements. + """ + pass + + def generate_env_name(self) -> str: + """ + Generates a unique environment name by hashing the model directory + and appending it to the base name. + + Returns: + str: A unique name for the environment. + """ + dir_hash = hashlib.md5(self.model_directory.encode()).hexdigest()[:8] + return f"{self.base_name}_{dir_hash}" + + +class CondaEnvironmentManager(EnvironmentManager): + """ + Manages a conda (or mamba) environment, providing methods to create, check, + and manipulate conda environments specifically. + """ + + def __init__(self, base_name: str, model_directory: str): + """ + Initializes the Conda environment manager with the specified base name + and model directory. It also generates the environment name and detects + the package manager (conda or mamba) to install dependencies.. + + Args: + base_name (str): The base name, i.e., model name, for the conda environment. + model_directory (str): The directory containing the model files. + """ + self.base_name = base_name + self.model_directory = model_directory + self.env_name = self.generate_env_name() + self.package_manager = self.detect_package_manager() + + @staticmethod + def detect_package_manager(): + """ + Detects whether 'mamba' or 'conda' is available as the package manager. + + Returns: + str: The name of the detected package manager ('mamba' or 'conda'). + """ + if shutil.which("mamba"): + log.info("Mamba detected, using mamba as package manager.") + return "mamba" + log.info("Mamba not detected, using conda as package manager.") + return "conda" + + def create_environment(self, force=False): + """ + Creates a conda environment using either an environment.yml file or + the specified Python version in setup.py/setup.cfg or project/toml. + If 'force' is True, any existing environment with the same name will + be removed first. + + Args: + force (bool): Whether to forcefully remove an existing environment. + """ + if force and self.env_exists(): + log.info(f"Removing existing conda environment: {self.env_name}") + subprocess.run( + [ + self.package_manager, + "env", + "remove", + "--name", + self.env_name, + "--yes", + ] + ) + + if not self.env_exists(): + env_file = os.path.join(self.model_directory, "environment.yml") + if os.path.exists(env_file): + log.info(f"Creating sub-conda environment {self.env_name} from environment.yml") + subprocess.run( + [ + self.package_manager, + "env", + "create", + "--name", + self.env_name, + "--file", + env_file, + ] + ) + else: + python_version = self.detect_python_version() + log.info( + f"Creating sub-conda environment {self.env_name} with Python {python_version}" + ) + subprocess.run( + [ + self.package_manager, + "create", + "--name", + self.env_name, + "--yes", + f"python={python_version}", + ] + ) + log.info(f"\tSub-conda environment created: {self.env_name}") + + self.install_dependencies() + + def env_exists(self) -> bool: + """ + Checks if the conda environment exists by querying the list of + existing conda environments. + + Returns: + bool: True if the conda environment exists, False otherwise. + """ + result = subprocess.run(["conda", "env", "list"], stdout=subprocess.PIPE) + return self.env_name in result.stdout.decode() + + def detect_python_version(self) -> str: + """ + Determines the required Python version from setup files in the model directory. + It checks 'setup.py', 'pyproject.toml', and 'setup.cfg' (in that order), for + version specifications. + + Returns: + str: The detected or default Python version. + """ + setup_py = os.path.join(self.model_directory, "setup.py") + pyproject_toml = os.path.join(self.model_directory, "pyproject.toml") + setup_cfg = os.path.join(self.model_directory, "setup.cfg") + current_python_version = ( + f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + ) + + def parse_version(version_str): + # Extract the first valid version number + import re + + match = re.search(r"\d+(.\d+)*", version_str) + return match.group(0) if match else current_python_version + + def is_version_compatible(requirement, current_version): + try: + specifier = SpecifierSet(requirement) + return current_version in specifier + except Exception as e: + log.error(f"Invalid specifier: {requirement}. Error: {e}") + return False + + if os.path.exists(setup_py): + with open(setup_py) as f: + for line in f: + if "python_requires" in line: + required_version = line.split("=")[1].strip() + if is_version_compatible(required_version, current_python_version): + log.info(f"Using current Python version: {current_python_version}") + return current_python_version + return parse_version(required_version) + + if os.path.exists(pyproject_toml): + with open(pyproject_toml) as f: + for line in f: + if "python" in line and "=" in line: + required_version = line.split("=")[1].strip() + if is_version_compatible(required_version, current_python_version): + log.info(f"Using current Python version: {current_python_version}") + return current_python_version + return parse_version(required_version) + + if os.path.exists(setup_cfg): + config = configparser.ConfigParser() + config.read(setup_cfg) + if "options" in config and "python_requires" in config["options"]: + required_version = config["options"]["python_requires"].strip() + if is_version_compatible(required_version, current_python_version): + log.info(f"Using current Python version: {current_python_version}") + return current_python_version + return parse_version(required_version) + + return current_python_version + + def install_dependencies(self): + """ + Installs dependencies in the conda environment using pip, based on the + model setup file + """ + log.info(f"Installing dependencies in conda environment: {self.env_name}") + cmd = [ + self.package_manager, + "run", + "-n", + self.env_name, + "pip", + "install", + "-e", + self.model_directory, + ] + subprocess.run(cmd, check=True) + + def run_command(self, command): + """ + Runs a specified command within the conda environment + Args: + command (str): The command to be executed in the conda environment. + """ + cmd = [ + "bash", + "-c", + f"{self.package_manager} run -n {self.env_name} {command}", + ] + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + ) + for line in process.stdout: + log.info(f"\t{line[:-1]}") + process.wait() + + +class VenvEnvironmentManager(EnvironmentManager): + """ + Manages a virtual environment created using Python's venv module. + Provides methods to create, check, and manipulate virtual environments. + """ + + def __init__(self, base_name: str, model_directory: str): + """ + Initializes the virtual environment manager with the specified base name + and model directory. + + Args: + base_name (str): The base name (i.e., model name) for the virtual environment. + model_directory (str): The directory containing the model files. + """ + + self.base_name = base_name + self.model_directory = model_directory + self.env_name = self.generate_env_name() + self.env_path = os.path.join(model_directory, self.env_name) + + def create_environment(self, force=False): + """ + Creates a virtual environment in the specified model directory. If 'force' + is True, any existing virtual environment will be removed before creation. + + Args: + force (bool): Whether to forcefully remove an existing virtual environment. + """ + if force and self.env_exists(): + log.info(f"Removing existing virtual environment: {self.env_name}") + shutil.rmtree(self.env_path) + + if not self.env_exists(): + log.info(f"Creating virtual environment: {self.env_name}") + venv.create(self.env_path, with_pip=True) + log.info(f"\tVirtual environment created: {self.env_name}") + self.install_dependencies() + + def env_exists(self) -> bool: + """ + Checks if the virtual environment exists by verifying the presence of its directory. + + Returns: + bool: True if the virtual environment exists, False otherwise. + """ + return os.path.isdir(self.env_path) + + def install_dependencies(self): + """ + Installs dependencies in the virtual environment using pip, based on the + model directory's configuration. + """ + log.info(f"Installing dependencies in virtual environment: {self.env_name}") + pip_executable = os.path.join(self.env_path, "bin", "pip") + cmd = f"{pip_executable} install -e {os.path.abspath(self.model_directory)}" + self.run_command(cmd) + + def run_command(self, command): + """ + Executes a specified command in the virtual environment and logs the output. + + Args: + command (str): The command to be executed in the virtual environment. + """ + env = os.environ.copy() + env.pop("PYTHONPATH", None) + process = subprocess.Popen( + command, + shell=True, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + ) + for line in process.stdout: + stripped_line = line.strip() + print(f"Logging: {stripped_line}") # Debug statement + log.info(stripped_line) + process.wait() + + +class DockerEnvironmentManager(EnvironmentManager): + """ + Manages a Docker environment, providing methods to create, check, + and manipulate Docker containers for the environment. + """ + + def __init__(self, base_name: str, model_directory: str): + self.base_name = base_name + self.model_directory = model_directory + + def create_environment(self, force=False): + pass + + def env_exists(self): + pass + + def run_command(self, command): + pass + + def install_dependencies(self): + pass + + +class EnvironmentFactory: + """ + Factory class for creating instances of environment managers based on the specified type. + """ + + @staticmethod + def get_env( + build: str = None, model_name: str = "model", model_path: str = None + ) -> EnvironmentManager: + """ + Returns an instance of an environment manager based on the specified build type. + It checks the current environment type and can return a conda, venv, or Docker + environment manager. + + Args: + build (str): The desired type of environment ('conda', 'venv', or 'docker'). + model_name (str): The name of the model for which the environment is being created. + model_path (str): The path to the model directory. + + Returns: + EnvironmentManager: An instance of the appropriate environment manager. + + Raises: + Exception: If an invalid environment type is specified. + """ + run_env = EnvironmentFactory.check_environment_type() + if run_env != build and build and build != "docker": + log.warning( + f"Selected build environment ({build}) for this model is different than that of" + f" the experiment run. Consider selecting the same environment." + ) + if build == "conda" or (not build and run_env == "conda"): + return CondaEnvironmentManager( + base_name=f"{model_name}", + model_directory=os.path.abspath(model_path), + ) + elif build == "venv" or (not build and run_env == "venv"): + return VenvEnvironmentManager( + base_name=f"{model_name}", + model_directory=os.path.abspath(model_path), + ) + elif build == "docker": + return DockerEnvironmentManager( + base_name=f"{model_name}", + model_directory=os.path.abspath(model_path), + ) + else: + raise Exception( + "Wrong environment selection. Please choose between " + '"conda", "venv" or "docker".' + ) + + @staticmethod + def check_environment_type(): + if "VIRTUAL_ENV" in os.environ: + return "venv" + try: + subprocess.run( + ["conda", "info"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + return "conda" + except FileNotFoundError: + pass + return None + + +if __name__ == "__main__": + + env = EnvironmentFactory.get_env( + "conda", model_path="../examples/case_h/models/pymock_poisson" + ) + env.create_environment(force=True) diff --git a/floatcsep/model.py b/floatcsep/model.py index 3d81bb8..aae3317 100644 --- a/floatcsep/model.py +++ b/floatcsep/model.py @@ -1,7 +1,13 @@ import json import logging +import sys import os import subprocess +import hashlib +import shutil +import venv +import configparser + from abc import ABC, abstractmethod from datetime import datetime from typing import List, Callable, Union, Mapping, Sequence @@ -124,9 +130,7 @@ def get_source( elif giturl: log.info(f"Retrieving model {self.name} from git url: " f"{giturl}") try: - from_git( - giturl, self.dir if self.path.fmt else self.path("path"), **kwargs - ) + from_git(giturl, self.dir if self.path.fmt else self.path("path"), **kwargs) except (git.NoSuchPathError, git.CommandError) as msg: raise git.NoSuchPathError(f"git url was not found {msg}") else: @@ -177,9 +181,7 @@ def iter_attr(val): return _get_value(val) list_walk = [ - (i, j) - for i, j in sorted(self.__dict__.items()) - if not i.startswith("_") and j + (i, j) for i, j in sorted(self.__dict__.items()) if not i.startswith("_") and j ] dict_walk = {i: j for i, j in list_walk} @@ -223,9 +225,7 @@ class TimeIndependentModel(Model): store_db (bool): flag to indicate whether to store the model in a database. """ - def __init__( - self, name: str, model_path: str, forecast_unit=1, store_db=False, **kwargs - ): + def __init__(self, name: str, model_path: str, forecast_unit=1, store_db=False, **kwargs): super().__init__(name, model_path, **kwargs) self.forecast_unit = forecast_unit self.store_db = store_db @@ -289,9 +289,7 @@ def stage(self, timewindows: Union[str, List[datetime]] = None) -> None: def get_forecast( self, tstring: Union[str, list] = None, region=None - ) -> Union[ - GriddedForecast, CatalogForecast, List[GriddedForecast], List[CatalogForecast] - ]: + ) -> Union[GriddedForecast, CatalogForecast, List[GriddedForecast], List[CatalogForecast]]: """ Wrapper that just returns a forecast when requested. """ @@ -333,9 +331,7 @@ def create_forecast(self, tstring: str, **kwargs) -> None: start_date, end_date = str2timewindow(tstring) self.forecast_from_file(start_date, end_date, **kwargs) - def forecast_from_file( - self, start_date: datetime, end_date: datetime, **kwargs - ) -> None: + def forecast_from_file(self, start_date: datetime, end_date: datetime, **kwargs) -> None: """ Generates a forecast from a file, by parsing and scaling it to. @@ -461,9 +457,7 @@ def stage(self, timewindows=None) -> None: def get_forecast( self, tstring: Union[str, list] = None, region=None - ) -> Union[ - GriddedForecast, CatalogForecast, List[GriddedForecast], List[CatalogForecast] - ]: + ) -> Union[GriddedForecast, CatalogForecast, List[GriddedForecast], List[CatalogForecast]]: """Wrapper that just returns a forecast, hiding the access method under the hood""" if isinstance(tstring, str): @@ -511,9 +505,7 @@ def create_forecast(self, tstring: str, **kwargs) -> None: else: log.info(f"Forecast of {tstring} of model {self.name} already " f"exists") - def forecast_from_func( - self, start_date: datetime, end_date: datetime, **kwargs - ) -> None: + def forecast_from_func(self, start_date: datetime, end_date: datetime, **kwargs) -> None: self.prepare_args(start_date, end_date, **kwargs) log.info( diff --git a/floatcsep/registry.py b/floatcsep/registry.py index 8ece9b2..b36c119 100644 --- a/floatcsep/registry.py +++ b/floatcsep/registry.py @@ -118,8 +118,7 @@ def build_tree( # set forecast names fc_files = { - win: join(dirtree["forecasts"], f"{prefix}_{win}.csv") - for win in windows + win: join(dirtree["forecasts"], f"{prefix}_{win}.csv") for win in windows } fc_exists = { diff --git a/pyproject.toml b/pyproject.toml index dc80b62..5eb4733 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,4 +6,9 @@ build-backend = "setuptools.build_meta" addopts = "--cov=floatcsep" testpaths = [ "tests", -] \ No newline at end of file +] + +[tool.black] +line-length = 96 +skip-string-normalization = false +target-version = ["py39", "py310", "py311"] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d1cbbb8..81d20f8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,8 @@ flake8 gitpython h5py matplotlib +packaging +pandas pycsep pyshp pyyaml diff --git a/requirements_dev.txt b/requirements_dev.txt index 9cc3cc8..258b89e 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -1,5 +1,6 @@ numpy cartopy +black dateparser docker flake8 @@ -10,9 +11,9 @@ matplotlib mercantile mypy obspy +packaging pandas pillow -pyblack pycsep pydocstringformatter pyproj diff --git a/setup.cfg b/setup.cfg index 427f38d..8eea110 100644 --- a/setup.cfg +++ b/setup.cfg @@ -27,6 +27,8 @@ install_requires = gitpython h5py matplotlib + packaging + pandas pycsep pyshp pyyaml @@ -55,6 +57,7 @@ dev = mercantile mypy obspy + packaging pandas pillow pycsep diff --git a/tests/artifacts/models/td_model/setup.cfg b/tests/artifacts/models/td_model/setup.cfg new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_environments.py b/tests/unit/test_environments.py new file mode 100644 index 0000000..e0013f7 --- /dev/null +++ b/tests/unit/test_environments.py @@ -0,0 +1,365 @@ +import venv +import unittest +import subprocess +import os +from unittest.mock import patch, MagicMock, call, mock_open +import shutil +import hashlib +import logging +from floatcsep.environments import ( + CondaEnvironmentManager, + EnvironmentFactory, + VenvEnvironmentManager, + DockerEnvironmentManager, +) + + +class TestCondaEnvironmentManager(unittest.TestCase): + + @classmethod + def setUpClass(cls): + if not shutil.which("conda"): + raise unittest.SkipTest("Conda is not available in the environment.") + + def setUp(self): + self.manager = CondaEnvironmentManager( + base_name="test_env", model_directory="/tmp/test_model" + ) + os.makedirs("/tmp/test_model", exist_ok=True) + with open("/tmp/test_model/environment.yml", "w") as f: + f.write("name: test_env\ndependencies:\n - python=3.8\n - numpy") + with open("/tmp/test_model/setup.py", "w") as f: + f.write( + "from setuptools import setup\nsetup(name='test_model', version='0.1')" + ) + + def tearDown(self): + if self.manager.env_exists(): + subprocess.run( + ["conda", "env", "remove", "--name", self.manager.env_name, "--yes"], + check=True, + ) + if os.path.exists("/tmp/test_model"): + shutil.rmtree("/tmp/test_model") + + @patch("subprocess.run") + @patch("shutil.which", return_value="conda") + def test_generate_env_name(self, mock_which, mock_run): + manager = CondaEnvironmentManager("test_base", "/path/to/model") + expected_name = ( + "test_base_" + hashlib.md5("/path/to/model".encode()).hexdigest()[:8] + ) + print(expected_name) + self.assertEqual(manager.generate_env_name(), expected_name) + + @patch("subprocess.run") + def test_env_exists(self, mock_run): + hashed = hashlib.md5("/path/to/model".encode()).hexdigest()[:8] + mock_run.return_value.stdout = f"test_base_{hashed}\n".encode() + + manager = CondaEnvironmentManager("test_base", "/path/to/model") + self.assertTrue(manager.env_exists()) + + @patch("subprocess.run") + @patch("os.path.exists", return_value=True) + def test_create_environment(self, mock_exists, mock_run): + manager = CondaEnvironmentManager("test_base", "/path/to/model") + manager.create_environment(force=False) + + expected_calls = [ + call(["conda", "env", "list"], stdout=-1), + call().stdout.decode(), + call().stdout.decode().__contains__(manager.env_name), + call( + [ + "mamba", + "env", + "create", + "--name", + manager.env_name, + "--file", + "/path/to/model/environment.yml", + ] + ), + call( + [ + "mamba", + "run", + "-n", + manager.env_name, + "pip", + "install", + "-e", + "/path/to/model", + ], + check=True, + ), + ] + + self.assertEqual(mock_run.call_count, 3) + mock_run.assert_has_calls(expected_calls, any_order=False) + + @patch("subprocess.run") + def test_create_environment_force(self, mock_run): + manager = CondaEnvironmentManager("test_base", "/path/to/model") + manager.env_exists = MagicMock(return_value=True) + manager.create_environment(force=True) + self.assertEqual(mock_run.call_count, 2) # One for remove, one for create + + @patch("subprocess.run") + @patch.object( + CondaEnvironmentManager, "detect_package_manager", return_value="conda" + ) + def test_install_dependencies(self, mock_detect_package_manager, mock_run): + manager = CondaEnvironmentManager("test_base", "/path/to/model") + manager.install_dependencies() + mock_run.assert_called_once_with( + [ + "conda", + "run", + "-n", + manager.env_name, + "pip", + "install", + "-e", + "/path/to/model", + ], + check=True, + ) + + @patch("shutil.which", return_value="conda") + @patch("os.path.exists", side_effect=[False, False, True]) + @patch( + "builtins.open", + new_callable=mock_open, + read_data="[metadata]\nname = test\n\n[options]\ninstall_requires =\n numpy\npython_requires = >=3.9,<3.12\n", + ) + def test_detect_python_version_setup_cfg(self, mock_open, mock_exists, mock_which): + manager = CondaEnvironmentManager("test_base", "../artifacts/models/td_model") + python_version = manager.detect_python_version() + + # Extract major and minor version parts + major_minor_version = ".".join(python_version.split(".")[:2]) + + self.assertIn( + major_minor_version, ["3.9", "3.10", "3.11"] + ) # Check if it falls within the specified range + + def test_create_and_delete_environment(self): + # Create the environment + self.manager.create_environment(force=True) + + # Check if the environment was created + result = subprocess.run( + ["conda", "env", "list"], stdout=subprocess.PIPE, check=True + ) + self.assertIn(self.manager.env_name, result.stdout.decode()) + + # Check if numpy is installed + result = subprocess.run( + [ + "conda", + "run", + "-n", + self.manager.env_name, + "python", + "-c", + "import numpy", + ], + check=True, + ) + self.assertEqual(result.returncode, 0) + + # Delete the environment + self.manager.create_environment( + force=True + ) # This should remove and recreate the environment + + # Check if the environment was recreated + result = subprocess.run( + ["conda", "env", "list"], stdout=subprocess.PIPE, check=True + ) + self.assertIn(self.manager.env_name, result.stdout.decode()) + + +class TestEnvironmentFactory(unittest.TestCase): + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value="conda") + def test_get_env_conda(self, mock_check_env, mock_abspath): + env_manager = EnvironmentFactory.get_env( + build="conda", model_name="test_model", model_path="/path/to/model" + ) + self.assertIsInstance(env_manager, CondaEnvironmentManager) + self.assertEqual(env_manager.base_name, "test_model") + self.assertEqual(env_manager.model_directory, "/absolute/path/to/model") + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value="venv") + def test_get_env_venv(self, mock_check_env, mock_abspath): + env_manager = EnvironmentFactory.get_env( + build="venv", model_name="test_model", model_path="/path/to/model" + ) + self.assertIsInstance(env_manager, VenvEnvironmentManager) + self.assertEqual(env_manager.base_name, "test_model") + self.assertEqual(env_manager.model_directory, "/absolute/path/to/model") + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value=None) + def test_get_env_docker(self, mock_check_env, mock_abspath): + env_manager = EnvironmentFactory.get_env( + build="docker", model_name="test_model", model_path="/path/to/model" + ) + self.assertIsInstance(env_manager, DockerEnvironmentManager) + self.assertEqual(env_manager.base_name, "test_model") + self.assertEqual(env_manager.model_directory, "/absolute/path/to/model") + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value="conda") + def test_get_env_default_conda(self, mock_check_env, mock_abspath): + env_manager = EnvironmentFactory.get_env( + build=None, model_name="test_model", model_path="/path/to/model" + ) + self.assertIsInstance(env_manager, CondaEnvironmentManager) + self.assertEqual(env_manager.base_name, "test_model") + self.assertEqual(env_manager.model_directory, "/absolute/path/to/model") + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value="venv") + def test_get_env_default_venv(self, mock_check_env, mock_abspath): + env_manager = EnvironmentFactory.get_env( + build=None, model_name="test_model", model_path="/path/to/model" + ) + self.assertIsInstance(env_manager, VenvEnvironmentManager) + self.assertEqual(env_manager.base_name, "test_model") + self.assertEqual(env_manager.model_directory, "/absolute/path/to/model") + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value=None) + def test_get_env_invalid(self, mock_check_env, mock_abspath): + with self.assertRaises(Exception) as context: + EnvironmentFactory.get_env( + build="invalid", model_name="test_model", model_path="/path/to/model" + ) + self.assertTrue("Wrong environment selection" in str(context.exception)) + + @patch("os.path.abspath", return_value="/absolute/path/to/model") + @patch.object(EnvironmentFactory, "check_environment_type", return_value="venv") + @patch("logging.Logger.warning") + def test_get_env_warning(self, mock_log_warning, mock_check_env, mock_abspath): + EnvironmentFactory.get_env( + build="conda", model_name="test_model", model_path="/path/to/model" + ) + mock_log_warning.assert_called_once_with( + f"Selected build environment (conda) for this model is different than that of" + f" the experiment run. Consider selecting the same environment." + ) + + +class TestVenvEnvironmentManager(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # Check if venv is available (Python standard library) + if not hasattr(venv, "create"): + raise unittest.SkipTest("Venv is not available in the environment.") + + def setUp(self): + self.model_directory = "/tmp/test_model" + self.manager = VenvEnvironmentManager( + base_name="test_env", model_directory=self.model_directory + ) + os.makedirs(self.model_directory, exist_ok=True) + with open(os.path.join(self.model_directory, "setup.py"), "w") as f: + f.write( + "from setuptools import setup\nsetup(name='test_model', version='0.1')" + ) + logging.disable(logging.CRITICAL) + + def tearDown(self): + if self.manager.env_exists(): + shutil.rmtree(self.manager.env_path) + if os.path.exists(self.model_directory): + shutil.rmtree(self.model_directory) + + def test_create_and_delete_environment(self): + # Create the environment + self.manager.create_environment(force=True) + + # Check if the environment was created + self.assertTrue(self.manager.env_exists()) + + # Check if pip is available in the environment + pip_executable = os.path.join(self.manager.env_path, "bin", "pip") + result = subprocess.run( + [pip_executable, "list"], stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + self.assertEqual(result.returncode, 0) # pip should run without errors + + # Delete the environment + self.manager.create_environment( + force=True + ) # This should remove and recreate the environment + + # Check if the environment was recreated + self.assertTrue(self.manager.env_exists()) + + def test_init(self): + self.assertEqual(self.manager.base_name, "test_env") + self.assertEqual(self.manager.model_directory, self.model_directory) + self.assertTrue(self.manager.env_name.startswith("test_env_")) + + def test_env_exists(self): + self.assertFalse(self.manager.env_exists()) + self.manager.create_environment(force=True) + self.assertTrue(self.manager.env_exists()) + + def test_create_environment(self): + self.manager.create_environment(force=True) + self.assertTrue(self.manager.env_exists()) + + def test_create_environment_force(self): + self.manager.create_environment(force=True) + env_path_before = self.manager.env_path + self.manager.create_environment(force=True) + self.assertTrue(self.manager.env_exists()) + self.assertEqual( + env_path_before, self.manager.env_path + ) # Ensure it's a new path + + def test_install_dependencies(self): + self.manager.create_environment(force=True) + pip_executable = os.path.join(self.manager.env_path, "bin", "pip") + result = subprocess.run( + [pip_executable, "install", "-e", self.model_directory], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + self.assertEqual(result.returncode, 0) # pip should run without errors + + @patch('subprocess.Popen') + def test_run_command(self, mock_popen): + # Arrange + mock_process = MagicMock() + mock_process.stdout = iter(["Output line 1\n", "Output line 2\n"]) + mock_process.wait.return_value = None + mock_popen.return_value = mock_process + + command = "echo test_command" + + # Act + self.manager.run_command(command) + + # Assert + mock_popen.assert_called_once_with( + command, + shell=True, + env=unittest.mock.ANY, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + ) + + +if __name__ == "__main__": + unittest.main()