From 08657be1a1e65148c960dbf6100e3a89ef350d84 Mon Sep 17 00:00:00 2001 From: Michal Hucko Date: Wed, 28 Feb 2024 16:30:05 +0100 Subject: [PATCH] Implement initialize (#29) Implements the `dss initialize` command to initialize dss on a microk8s cluster. The command creates all resources needed for the basic dss experience (namespace, mlflow, etc) --- .github/workflows/tests.yaml | 57 ++++++++++---- requirements-fmt.txt | 6 +- requirements-lint.in | 1 + requirements-lint.txt | 47 ++++++++++-- requirements-unit.in | 1 + requirements-unit.txt | 133 +++++++++++++++++++++++++++++++-- requirements.in | 4 + requirements.txt | 82 ++++++++++++++++++++ setup.cfg | 6 ++ src/dss/initialize.py | 57 ++++++++++++++ src/dss/logger.py | 47 ++++++++++++ src/dss/main.py | 34 +++++++++ src/dss/manifests.yaml | 75 +++++++++++++++++++ src/dss/utils.py | 83 ++++++++++++++++++++ tests/integration/test_dss.py | 68 +++++++++++++++++ tests/unit/prepare_host_env.py | 1 - tests/unit/test_initialize.py | 73 ++++++++++++++++++ tests/unit/test_utils.py | 121 ++++++++++++++++++++++++++++++ tox.ini | 19 ++++- 19 files changed, 880 insertions(+), 35 deletions(-) create mode 100644 requirements.in create mode 100644 requirements.txt create mode 100644 src/dss/initialize.py create mode 100644 src/dss/logger.py create mode 100644 src/dss/main.py create mode 100644 src/dss/manifests.yaml create mode 100644 src/dss/utils.py create mode 100644 tests/integration/test_dss.py delete mode 100644 tests/unit/prepare_host_env.py create mode 100644 tests/unit/test_initialize.py create mode 100644 tests/unit/test_utils.py diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 89b4a87..79a3f43 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -1,26 +1,55 @@ name: Tests - on: workflow_call: - + jobs: lint: name: Lint runs-on: ubuntu-22.04 steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Install dependencies - run: python3 -m pip install tox - - name: Run linters - run: tox -e lint + - name: Checkout repository + uses: actions/checkout@v4 + - name: Install dependencies + run: python3 -m pip install tox + - name: Run linters + run: tox -e lint + unit-test: name: Unit tests runs-on: ubuntu-22.04 steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Install dependencies - run: python -m pip install tox - - name: Run tests - run: tox -e unit + - name: Checkout repository + uses: actions/checkout@v4 + - name: Install dependencies + run: python -m pip install tox + - name: Run tests + run: tox -e unit + + integration: + name: Integration tests + runs-on: ubuntu-22.04 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + # Remove once https://github.com/canonical/bundle-kubeflow/issues/761 + # is resolved and applied to all ROCKs repositories + - name: Install python version from input + run: | + sudo add-apt-repository ppa:deadsnakes/ppa -y + sudo apt update -y + VERSION=3.10 + sudo apt install python3.10 python3.10-distutils python3.10-venv -y + wget https://bootstrap.pypa.io/get-pip.py + python3.10 get-pip.py + python3.10 -m pip install tox + rm get-pip.py + + - uses: balchua/microk8s-actions@v0.3.2 + with: + channel: '1.28/stable' + addons: '["hostpath-storage"]' + + - name: Install library + run: sg microk8s -c "tox -vve integration" + diff --git a/requirements-fmt.txt b/requirements-fmt.txt index 5eaa044..7bd6d98 100644 --- a/requirements-fmt.txt +++ b/requirements-fmt.txt @@ -2,14 +2,14 @@ # This file is autogenerated by pip-compile with Python 3.10 # by the following command: # -# pip-compile ./requirements-fmt.in +# pip-compile requirements-fmt.in # black==23.12.1 - # via -r ./requirements-fmt.in + # via -r requirements-fmt.in click==8.1.7 # via black isort==5.13.2 - # via -r ./requirements-fmt.in + # via -r requirements-fmt.in mypy-extensions==1.0.0 # via black packaging==23.2 diff --git a/requirements-lint.in b/requirements-lint.in index c453874..87f005e 100644 --- a/requirements-lint.in +++ b/requirements-lint.in @@ -4,3 +4,4 @@ flake8-builtins flake8-copyright pep8-naming pyproject-flake8 +-r requirements-fmt.txt diff --git a/requirements-lint.txt b/requirements-lint.txt index 39a5fbe..9b84d64 100644 --- a/requirements-lint.txt +++ b/requirements-lint.txt @@ -2,32 +2,63 @@ # This file is autogenerated by pip-compile with Python 3.10 # by the following command: # -# pip-compile ./requirements-lint.in +# pip-compile requirements-lint.in # +black==23.12.1 + # via -r requirements-fmt.txt +click==8.1.7 + # via + # -r requirements-fmt.txt + # black codespell==2.2.6 - # via -r ./requirements-lint.in + # via -r requirements-lint.in flake8==6.1.0 # via - # -r ./requirements-lint.in + # -r requirements-lint.in # flake8-builtins # pep8-naming # pyproject-flake8 flake8-builtins==2.2.0 - # via -r ./requirements-lint.in + # via -r requirements-lint.in flake8-copyright==0.2.4 - # via -r ./requirements-lint.in + # via -r requirements-lint.in +isort==5.13.2 + # via -r requirements-fmt.txt mccabe==0.7.0 # via flake8 +mypy-extensions==1.0.0 + # via + # -r requirements-fmt.txt + # black +packaging==23.2 + # via + # -r requirements-fmt.txt + # black +pathspec==0.12.1 + # via + # -r requirements-fmt.txt + # black pep8-naming==0.13.3 - # via -r ./requirements-lint.in + # via -r requirements-lint.in +platformdirs==4.1.0 + # via + # -r requirements-fmt.txt + # black pycodestyle==2.11.1 # via flake8 pyflakes==3.1.0 # via flake8 pyproject-flake8==6.1.0 - # via -r ./requirements-lint.in + # via -r requirements-lint.in tomli==2.0.1 - # via pyproject-flake8 + # via + # -r requirements-fmt.txt + # black + # pyproject-flake8 +typing-extensions==4.9.0 + # via + # -r requirements-fmt.txt + # black # The following packages are considered to be unsafe in a requirements file: # setuptools diff --git a/requirements-unit.in b/requirements-unit.in index 89b68f8..1d2112e 100644 --- a/requirements-unit.in +++ b/requirements-unit.in @@ -1,3 +1,4 @@ pytest pytest-mock coverage[toml] +-r requirements.txt diff --git a/requirements-unit.txt b/requirements-unit.txt index c08bec2..b403cda 100644 --- a/requirements-unit.txt +++ b/requirements-unit.txt @@ -2,25 +2,148 @@ # This file is autogenerated by pip-compile with Python 3.10 # by the following command: # -# pip-compile ./requirements-unit.in +# pip-compile requirements-unit.in # +anyio==4.2.0 + # via + # -r requirements.txt + # httpx +attrs==23.2.0 + # via + # -r requirements.txt + # jsonschema +certifi==2024.2.2 + # via + # -r requirements.txt + # httpcore + # httpx + # requests +charmed-kubeflow-chisme==0.2.1 + # via -r requirements.txt +charset-normalizer==3.3.2 + # via + # -r requirements.txt + # requests +click==8.1.7 + # via -r requirements.txt coverage[toml]==7.4.0 - # via -r ./requirements-unit.in + # via -r requirements-unit.in +deepdiff==6.2.1 + # via + # -r requirements.txt + # charmed-kubeflow-chisme exceptiongroup==1.2.0 - # via pytest + # via + # -r requirements.txt + # anyio + # pytest +h11==0.14.0 + # via + # -r requirements.txt + # httpcore +httpcore==1.0.2 + # via + # -r requirements.txt + # httpx +httpx==0.26.0 + # via + # -r requirements.txt + # lightkube +idna==3.6 + # via + # -r requirements.txt + # anyio + # httpx + # requests iniconfig==2.0.0 # via pytest +jinja2==3.1.3 + # via + # -r requirements.txt + # charmed-kubeflow-chisme +jsonschema==4.17.3 + # via + # -r requirements.txt + # serialized-data-interface +lightkube==0.15.1 + # via + # -r requirements.txt + # charmed-kubeflow-chisme +lightkube-models==1.29.0.7 + # via + # -r requirements.txt + # lightkube +markupsafe==2.1.5 + # via + # -r requirements.txt + # jinja2 +ops==2.10.0 + # via + # -r requirements.txt + # charmed-kubeflow-chisme + # serialized-data-interface +ordered-set==4.1.0 + # via + # -r requirements.txt + # deepdiff packaging==23.2 # via pytest pluggy==1.3.0 # via pytest +pyrsistent==0.20.0 + # via + # -r requirements.txt + # jsonschema pytest==7.4.4 # via - # -r ./requirements-unit.in + # -r requirements-unit.in # pytest-mock pytest-mock==3.12.0 - # via -r ./requirements-unit.in + # via -r requirements-unit.in +pyyaml==6.0.1 + # via + # -r requirements.txt + # lightkube + # ops + # serialized-data-interface +requests==2.31.0 + # via + # -r requirements.txt + # serialized-data-interface +ruamel-yaml==0.18.6 + # via + # -r requirements.txt + # charmed-kubeflow-chisme +ruamel-yaml-clib==0.2.8 + # via + # -r requirements.txt + # ruamel-yaml +serialized-data-interface==0.7.0 + # via + # -r requirements.txt + # charmed-kubeflow-chisme +sniffio==1.3.0 + # via + # -r requirements.txt + # anyio + # httpx +tenacity==8.2.3 + # via + # -r requirements.txt + # charmed-kubeflow-chisme tomli==2.0.1 # via # coverage # pytest +typing-extensions==4.9.0 + # via + # -r requirements.txt + # anyio +urllib3==2.2.0 + # via + # -r requirements.txt + # requests +websocket-client==1.7.0 + # via + # -r requirements.txt + # ops diff --git a/requirements.in b/requirements.in new file mode 100644 index 0000000..346a360 --- /dev/null +++ b/requirements.in @@ -0,0 +1,4 @@ +Click +charmed_kubeflow_chisme +lightkube +pyyaml diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c5da55d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,82 @@ +# +# This file is autogenerated by pip-compile with Python 3.10 +# by the following command: +# +# pip-compile requirements.in +# +anyio==4.2.0 + # via httpx +attrs==23.2.0 + # via jsonschema +certifi==2024.2.2 + # via + # httpcore + # httpx + # requests +charmed-kubeflow-chisme==0.2.1 + # via -r requirements.in +charset-normalizer==3.3.2 + # via requests +click==8.1.7 + # via -r requirements.in +deepdiff==6.2.1 + # via charmed-kubeflow-chisme +exceptiongroup==1.2.0 + # via anyio +h11==0.14.0 + # via httpcore +httpcore==1.0.2 + # via httpx +httpx==0.26.0 + # via lightkube +idna==3.6 + # via + # anyio + # httpx + # requests +jinja2==3.1.3 + # via charmed-kubeflow-chisme +jsonschema==4.17.3 + # via serialized-data-interface +lightkube==0.15.1 + # via + # -r requirements.in + # charmed-kubeflow-chisme +lightkube-models==1.29.0.7 + # via lightkube +markupsafe==2.1.5 + # via jinja2 +ops==2.10.0 + # via + # charmed-kubeflow-chisme + # serialized-data-interface +ordered-set==4.1.0 + # via deepdiff +pyrsistent==0.20.0 + # via jsonschema +pyyaml==6.0.1 + # via + # -r requirements.in + # lightkube + # ops + # serialized-data-interface +requests==2.31.0 + # via serialized-data-interface +ruamel-yaml==0.18.6 + # via charmed-kubeflow-chisme +ruamel-yaml-clib==0.2.8 + # via ruamel-yaml +serialized-data-interface==0.7.0 + # via charmed-kubeflow-chisme +sniffio==1.3.0 + # via + # anyio + # httpx +tenacity==8.2.3 + # via charmed-kubeflow-chisme +typing-extensions==4.9.0 + # via anyio +urllib3==2.2.0 + # via requests +websocket-client==1.7.0 + # via ops diff --git a/setup.cfg b/setup.cfg index ed409f5..282dcd4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -7,7 +7,13 @@ package_dir = = src packages = find: install_requires = + charmed-kubeflow-chisme Click + lightkube +include_package_data = True + +[options.package_data] +dss = manifests.yaml [options.packages.find] where = src diff --git a/src/dss/initialize.py b/src/dss/initialize.py new file mode 100644 index 0000000..4bb2b3c --- /dev/null +++ b/src/dss/initialize.py @@ -0,0 +1,57 @@ +import os + +from charmed_kubeflow_chisme.kubernetes import KubernetesResourceHandler +from lightkube import Client +from lightkube.resources.apps_v1 import Deployment +from lightkube.resources.core_v1 import Namespace, PersistentVolumeClaim, Service + +from dss.logger import setup_logger +from dss.utils import wait_for_deployment_ready + +# Set up logger +logger = setup_logger("logs/dss.log") + + +DSS_CLI_MANAGER_LABELS = {"app.kubernetes.io/managed-by": "dss-cli"} + + +def initialize(lightkube_client: Client) -> None: + """ + Initializes the Kubernetes cluster by applying manifests from a YAML file. + + Args: + lightkube_client (Client): The Kubernetes client. + + Returns: + None + """ + # Path to the manifests YAML file + manifests_file = os.path.join(os.path.dirname(__file__), "manifests.yaml") + + # Initialize KubernetesResourceHandler + k8s_resource_handler = KubernetesResourceHandler( + field_manager="dss", + labels=DSS_CLI_MANAGER_LABELS, + template_files=[manifests_file], + context={}, + resource_types={Deployment, Service, PersistentVolumeClaim, Namespace}, + lightkube_client=lightkube_client, + ) + + try: + # Apply resources using KubernetesResourceHandler + k8s_resource_handler.apply() + + # Wait for mlflow deployment to be ready + wait_for_deployment_ready(lightkube_client, namespace="dss", deployment_name="mlflow") + + logger.info( + "DSS initialized. To create your first notebook run the command:\n\ndss create-notebook" # noqa E501 + ) + + except TimeoutError: + logger.error( + "Timeout waiting for deployment 'mlflow-deployment' in namespace 'dss' to be ready. " # noqa E501 + "Deleting resources..." + ) + k8s_resource_handler.delete() diff --git a/src/dss/logger.py b/src/dss/logger.py new file mode 100644 index 0000000..dc5382a --- /dev/null +++ b/src/dss/logger.py @@ -0,0 +1,47 @@ +import logging +import os +from logging.handlers import RotatingFileHandler + + +def setup_logger(log_file_path: str, log_level: int = logging.DEBUG) -> logging.Logger: + """ + Set up a logger with a file handler and console handler. + + Args: + log_file_path (str): Path to the log file. + log_level (int, optional): Logging level. Defaults to logging.DEBUG. + + Returns: + logging.Logger: Configured logger object. + """ + logger = logging.getLogger(__name__) + + # Check if the logger already has handlers to avoid duplication + if not logger.handlers: + logger.setLevel(log_level) + + # Create log formatter + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] [%(module)s] [%(funcName)s]: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + # Create the log directory if it doesn't exist + if not os.path.exists(os.path.dirname(log_file_path)): + os.makedirs(os.path.dirname(log_file_path)) + + # Create file handler + file_handler = RotatingFileHandler(log_file_path, maxBytes=5 * 1024 * 1024, backupCount=5) + file_handler.setLevel(log_level) + file_handler.setFormatter(formatter) + + # Create console handler + console_handler = logging.StreamHandler() + console_handler.setLevel(log_level) + console_handler.setFormatter(formatter) + + # Add handlers to the logger + logger.addHandler(file_handler) + logger.addHandler(console_handler) + + return logger diff --git a/src/dss/main.py b/src/dss/main.py new file mode 100644 index 0000000..73d5766 --- /dev/null +++ b/src/dss/main.py @@ -0,0 +1,34 @@ +import click + +from dss.initialize import initialize +from dss.logger import setup_logger +from dss.utils import KUBECONFIG_DEFAULT, get_default_kubeconfig, get_lightkube_client + +# Set up logger +logger = setup_logger("logs/dss.log") + + +@click.group() +def main(): + """Command line interface for managing the DSS application.""" + + +@main.command(name="initialize") +@click.option( + "--kubeconfig", + help=f"Path to a Kubernetes config file. Defaults to the value of the KUBECONFIG environment variable, else to '{KUBECONFIG_DEFAULT}'.", # noqa E501 +) +def initialize_command(kubeconfig: str) -> None: + """ + Initialize DSS on the given Kubernetes cluster. + """ + logger.info("Executing initialize command") + + kubeconfig = get_default_kubeconfig(kubeconfig) + lightkube_client = get_lightkube_client(kubeconfig) + + initialize(lightkube_client=lightkube_client) + + +if __name__ == "__main__": + main() diff --git a/src/dss/manifests.yaml b/src/dss/manifests.yaml new file mode 100644 index 0000000..892af61 --- /dev/null +++ b/src/dss/manifests.yaml @@ -0,0 +1,75 @@ +--- +apiVersion: v1 +kind: Namespace +metadata: + name: dss + labels: + app.kubernetes.io/part-of: dss + +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: mlflow + namespace: dss + labels: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 1Gi + +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: mlflow + namespace: dss + labels: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss + template: + metadata: + labels: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss + spec: + containers: + - name: mlflow + image: ubuntu/mlflow:2.1.1_1.0-22.04 + args: ["mlflow", "server", "--host", "0.0.0.0", "--port", "5000"] + ports: + - containerPort: 5000 + volumeMounts: + - name: mlflow + mountPath: /mlruns + volumes: + - name: mlflow + persistentVolumeClaim: + claimName: mlflow + +--- +apiVersion: v1 +kind: Service +metadata: + name: mlflow + namespace: dss + labels: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss +spec: + selector: + app.kubernetes.io/name: mlflow + app.kubernetes.io/part-of: dss + ports: + - protocol: TCP + port: 5000 diff --git a/src/dss/utils.py b/src/dss/utils.py new file mode 100644 index 0000000..e744260 --- /dev/null +++ b/src/dss/utils.py @@ -0,0 +1,83 @@ +import os +import time +from typing import Optional + +from lightkube import Client, KubeConfig +from lightkube.resources.apps_v1 import Deployment + +from dss.logger import setup_logger + +# Set up logger +logger = setup_logger("logs/dss.log") + +# Name for the environment variable storing kubeconfig +KUBECONFIG_ENV_VAR = "DSS_KUBECONFIG" +KUBECONFIG_DEFAULT = "./kubeconfig" + + +def wait_for_deployment_ready( + client: Client, + namespace: str, + deployment_name: str, + timeout_seconds: int = 60, + interval_seconds: int = 10, +) -> None: + """ + Waits for a Kubernetes deployment to be ready. + + Args: + client (Client): The Kubernetes client. + namespace (str): The namespace of the deployment. + deployment_name (str): The name of the deployment. + timeout_seconds (int): Timeout in seconds. Defaults to 600. + interval_seconds (int): Interval between checks in seconds. Defaults to 10. + + Returns: + None + """ + logger.info( + f"Waiting for deployment {deployment_name} in namespace {namespace} to be ready..." + ) + start_time = time.time() + while True: + deployment: Deployment = client.get(Deployment, namespace=namespace, name=deployment_name) + if deployment.status.availableReplicas == deployment.spec.replicas: + logger.info(f"Deployment {deployment_name} in namespace {namespace} is ready") + break + elif time.time() - start_time >= timeout_seconds: + raise TimeoutError( + f"Timeout waiting for deployment {deployment_name} in namespace {namespace} to be ready" # noqa E501 + ) + else: + time.sleep(interval_seconds) + logger.info( + f"Waiting for deployment {deployment_name} in namespace {namespace} to be ready..." + ) + + +def get_default_kubeconfig(kubeconfig: Optional[str] = None) -> str: + """ + Get the kubeconfig file path, from input, environment variable, or default. + + Args: + kubeconfig (str): Path to a kubeconfig file. Defaults to None. + + Returns: + str: the value of kubeconfig if it is not None, else: + the value of the DSS_KUBECONFIG environment variable if it is set, else: + the default value of "./kubeconfig" + + """ + if kubeconfig: + return kubeconfig + elif os.environ.get(KUBECONFIG_ENV_VAR, ""): + return os.environ.get(KUBECONFIG_ENV_VAR, "") + else: + return KUBECONFIG_DEFAULT + + +def get_lightkube_client(kubeconfig: Optional[str] = None): + # Compute the default kubeconfig, if not specified + kubeconfig = KubeConfig.from_file(kubeconfig) + lightkube_client = Client(config=kubeconfig) + return lightkube_client diff --git a/tests/integration/test_dss.py b/tests/integration/test_dss.py new file mode 100644 index 0000000..168d76d --- /dev/null +++ b/tests/integration/test_dss.py @@ -0,0 +1,68 @@ +import subprocess + +import pytest +from charmed_kubeflow_chisme.kubernetes import KubernetesResourceHandler +from lightkube.resources.apps_v1 import Deployment +from lightkube.resources.core_v1 import Namespace, PersistentVolumeClaim, Service + +from dss.initialize import DSS_CLI_MANAGER_LABELS + + +def test_initialize_creates_dss(cleanup_after_initialize) -> None: + """ + Integration test to verify if the initialize command creates the 'dss' namespace and + the 'mlflow' deployment is active in the 'dss' namespace. + """ + # TODO: is there a better way to initialize this? Maybe an optional argument to the test? + kubeconfig = "~/.kube/config" + + result = subprocess.run( + ["dss", "initialize", "--kubeconfig", kubeconfig], + capture_output=True, + text=True, + ) + + # Check if the command executed successfully + print(result.stdout) + assert result.returncode == 0 + + # Check if the dss namespace exists using kubectl + kubectl_result = subprocess.run( + ["kubectl", "get", "namespace", "dss"], capture_output=True, text=True + ) + assert "dss" in kubectl_result.stdout + + # Check if the mlflow-deployment deployment is active in the dss namespace + kubectl_result = subprocess.run( + ["kubectl", "get", "deployment", "mlflow", "-n", "dss"], + capture_output=True, + text=True, + ) + assert "mlflow" in kubectl_result.stdout + assert ( + "1/1" in kubectl_result.stdout + ) # Assuming it should have 1 replica and all are available + + +@pytest.fixture() +def cleanup_after_initialize(): + """Cleans up resources that might have been deployed by dss initialize. + + Note that this is a white-box implementation - it depends on knowing what could be deployed and + explicitly removing those objects, rather than truly restoring the cluster to a previous state. + This could be leaky, depending on how `dss initialize` is changed in future. + """ + yield + + k8s_resource_handler = KubernetesResourceHandler( + field_manager="dss", + labels=DSS_CLI_MANAGER_LABELS, + template_files=[], + context={}, + resource_types={Deployment, Service, PersistentVolumeClaim, Namespace}, + ) + + # Attempt to clean up anything that initialize might create + # Note that .delete() does not wait on the objects to be successfully deleted, so repeating + # the tests quickly can still cause an issue + k8s_resource_handler.delete() diff --git a/tests/unit/prepare_host_env.py b/tests/unit/prepare_host_env.py deleted file mode 100644 index 0d4b760..0000000 --- a/tests/unit/prepare_host_env.py +++ /dev/null @@ -1 +0,0 @@ -# Test file \ No newline at end of file diff --git a/tests/unit/test_initialize.py b/tests/unit/test_initialize.py new file mode 100644 index 0000000..1c94c97 --- /dev/null +++ b/tests/unit/test_initialize.py @@ -0,0 +1,73 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from dss.initialize import initialize + + +@pytest.fixture +def mock_environ_get() -> MagicMock: + """ + Fixture to mock the os.environ.get function. + """ + with patch("dss.initialize.os.environ.get") as mock_env_get: + yield mock_env_get + + +@pytest.fixture +def mock_client() -> MagicMock: + """ + Fixture to mock the Client class. + """ + with patch("dss.initialize.Client") as mock_client: + yield mock_client + + +@pytest.fixture +def mock_resource_handler() -> MagicMock: + """ + Fixture to mock the KubernetesResourceHandler class. + """ + with patch("dss.initialize.KubernetesResourceHandler") as mock_handler: + yield mock_handler + + +@pytest.fixture +def mock_logger() -> MagicMock: + """ + Fixture to mock the logger object. + """ + with patch("dss.initialize.logger") as mock_logger: + yield mock_logger + + +def test_initialize_success( + mock_environ_get: MagicMock, + mock_client: MagicMock, + mock_resource_handler: MagicMock, + mock_logger: MagicMock, +) -> None: + """ + Test case to verify successful initialization. + """ + # Mock the behavior of Client + mock_client_instance = MagicMock() + mock_client.return_value = mock_client_instance + + # Mock the behavior of KubernetesResourceHandler + mock_resource_handler_instance = MagicMock() + mock_resource_handler.return_value = mock_resource_handler_instance + + # Mock wait_for_deployment_ready + with patch("dss.initialize.wait_for_deployment_ready") as mock_wait_for_deployment_ready: + # Call the function to test + initialize(lightkube_client=mock_client_instance) + + # Assertions + mock_resource_handler_instance.apply.assert_called_once() + mock_wait_for_deployment_ready.assert_called_once_with( + mock_client_instance, namespace="dss", deployment_name="mlflow" + ) + mock_logger.info.assert_called_with( + "DSS initialized. To create your first notebook run the command:\n\ndss create-notebook" + ) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py new file mode 100644 index 0000000..15f66e8 --- /dev/null +++ b/tests/unit/test_utils.py @@ -0,0 +1,121 @@ +from unittest.mock import MagicMock, patch + +import pytest +from lightkube.resources.apps_v1 import Deployment + +from dss.utils import ( + KUBECONFIG_DEFAULT, + get_default_kubeconfig, + get_lightkube_client, + wait_for_deployment_ready, +) + + +@pytest.fixture +def mock_client() -> MagicMock: + """ + Fixture to mock the Client class. + """ + with patch("dss.utils.Client") as mock_client: + yield mock_client + + +@pytest.fixture +def mock_environ_get() -> MagicMock: + """ + Fixture to mock the os.environ.get function. + """ + with patch("dss.utils.os.environ.get") as mock_env_get: + yield mock_env_get + + +@pytest.fixture +def mock_kubeconfig() -> MagicMock: + """ + Fixture to mock the KubeConfig.from_dict function. + """ + with patch("dss.utils.KubeConfig") as mock_kubeconfig: + yield mock_kubeconfig + + +@pytest.fixture +def mock_logger() -> MagicMock: + """ + Fixture to mock the logger object. + """ + with patch("dss.initialize.logger") as mock_logger: + yield mock_logger + + +@pytest.mark.parametrize( + "kubeconfig, kubeconfig_env_var, expected", + [ + ("some_file", "", "some_file"), + (None, "some_file", "some_file"), + (None, "", KUBECONFIG_DEFAULT), + ], +) +def test_get_default_kubeconfig_successful( + kubeconfig: str, + kubeconfig_env_var: str, + expected: str, + mock_environ_get: MagicMock, +) -> None: + """ + Test case to verify missing kubeconfig environment variable. + + Args: + kubeconfig: path to a kubeconfig file, passed to get_lightkube_client by arg + kubeconfig_env_var: environment variable for kubeconfig + expected: expected returned value for kubeconfig + """ + mock_environ_get.return_value = kubeconfig_env_var + + returned = get_default_kubeconfig(kubeconfig) + assert returned == expected + + +def test_get_lightkube_client_successful( + mock_kubeconfig: MagicMock, + mock_client: MagicMock, +) -> None: + """ + Tests that we successfully try to create a lightkube client, given a kubeconfig. + """ + kubeconfig = "some_file" + mock_kubeconfig_instance = "kubeconfig-returned" + mock_kubeconfig.from_file.return_value = mock_kubeconfig_instance + + returned_client = get_lightkube_client(kubeconfig) + + mock_kubeconfig.from_file.assert_called_with(kubeconfig) + mock_client.assert_called_with(config=mock_kubeconfig_instance) + assert returned_client is not None + + +def test_wait_for_deployment_ready_timeout(mock_client: MagicMock, mock_logger: MagicMock) -> None: + """ + Test case to verify timeout while waiting for deployment to be ready. + """ + # Mock the behavior of the client.get method to return a deployment with available replicas = 0 + mock_client_instance = MagicMock() + mock_client_instance.get.return_value = MagicMock( + spec=Deployment, status=MagicMock(availableReplicas=0), spec_replicas=1 + ) + + # Call the function to test + with pytest.raises(TimeoutError) as exc_info: + wait_for_deployment_ready( + mock_client_instance, + namespace="test-namespace", + deployment_name="test-deployment", + timeout_seconds=5, + interval_seconds=1, + ) + + # Assertions + assert ( + str(exc_info.value) + == "Timeout waiting for deployment test-deployment in namespace test-namespace to be ready" + ) + assert mock_client_instance.get.call_count == 6 # 5 attempts, 1 final attempt after timeout diff --git a/tox.ini b/tox.ini index 3a7bb37..5fc571b 100644 --- a/tox.ini +++ b/tox.ini @@ -7,7 +7,7 @@ max-line-length = 100 [tox] skipsdist=True skip_missing_interpreters = True -envlist = fmt, lint, unit, update-requirements +envlist = fmt, integration, lint, unit, update-requirements [vars] src_path = {toxinidir}/src/ @@ -28,9 +28,12 @@ allowlist_externals = find pip-compile xargs -commands = -; uses 'bash -c' because piping didn't work in regular tox commands - bash -c 'find . -type f -name "requirements*.in" | xargs --replace=\{\} pip-compile --resolver=backtracking \{\}' +commands = + ; we must preserve the order of compilation, since each *.in file depends on some *.txt file. + ; For example, requirements-unit.in depends on requirements.txt and we must compile first + ; requirements.txt to ensure that requirements-unit.txt get the same dependency as the requirements.txt + bash -c 'for pattern in "requirements.in" "requirements-fmt.in" "requirements*.in"; do find . -type f -name "$pattern" -exec bash -c "cd \$(dirname "{}") && pip-compile --resolver=backtracking \$(basename "{}")" \;; done' + deps = pip-tools description = Update requirements files by executing pip-compile on all requirements*.in files, including those in subdirs. @@ -69,3 +72,11 @@ commands = deps = -r requirements-unit.txt description = Run unit tests + +[testenv:integration] +commands = + pip install . + pytest {[vars]tst_path}/integration -v --tb native -s {posargs} +deps = + -r requirements-unit.txt +description = Run integration tests