diff --git a/.github/workflows/python-publish-dev.yml b/.github/workflows/python-publish-dev.yml new file mode 100644 index 0000000..693ae9d --- /dev/null +++ b/.github/workflows/python-publish-dev.yml @@ -0,0 +1,50 @@ +name: Upload dev Python Package + +# Publish a dev pre-release to PyPI on tag push (from any branch), so a change can be +# validated on a canary environment before merging. Stable releases are NOT published this way: +# they go through a published GitHub Release (python-publish.yml). +# +# Usage: git tag v7.0.0-dev.1 && git push origin v7.0.0-dev.1 +# Tags follow semantic versioning (vMAJOR.MINOR.PATCH-dev.N); the workflow translates to the +# PEP 440 equivalent (7.0.0.dev1) for the package. The version is derived from the tag (not +# from the VERSION file), and the workflow refuses anything that is not a vX.Y.Z-dev.N tag. +# pip only installs dev versions when pinned explicitly (or with --pre), so consumers can't +# pick one up by accident. + +on: + push: + tags: + - "v*-dev.*" + +jobs: + deploy: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - name: Refuse anything but a SemVer dev version tag + run: | + echo "${GITHUB_REF_NAME}" | grep -Eq '^v[0-9]+\.[0-9]+\.[0-9]+-dev\.[0-9]+$' || { + echo "Tag ${GITHUB_REF_NAME} is not a vX.Y.Z-dev.N dev version, refusing to publish." + exit 1 + } + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + - name: Set package version from the tag (SemVer -> PEP 440) + run: | + SEMVER="${GITHUB_REF_NAME#v}" + echo "${SEMVER/-dev./.dev}" > VERSION + echo "Publishing version $(cat VERSION) from tag ${GITHUB_REF_NAME}" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install build + - name: Build package + run: python -m build + - name: Publish package + uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f684eb..03d869e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,32 @@ # Changelog -## Unreleased +## v7.0.0 (2026-07-08) ### BREAKING CHANGE - Raise minimum supported Python version to 3.13. +### BREAKING CHANGE +- Drop the `wiremind-kubernetes` dependency (library being decommissioned): the subset + chartreuse uses (`KubernetesDeploymentManager` with `stop_pods`/`start_pods`, kube config + loading, `run_command`) now lives in `chartreuse.kubernetes_helper` and + `chartreuse.utils.command`, depending only on the official `kubernetes` client. + Library users importing chartreuse should not be affected unless they relied on chartreuse + pulling in wiremind-kubernetes transitively. + +### Feat +- New `chartreuse-restore` entrypoint, to run as a post-deployment (e.g. ArgoCD PostSync) hook + when `CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT` is enabled. On that path `chartreuse-upgrade` + skips `start_pods()` and relies on the following deployment to restore replicas, which never + happens for HPA-managed Deployments (their chart omits `spec.replicas`, and an HPA whose + target is at 0 replicas with `minReplicas >= 1` is ScalingDisabled): they stayed stranded at + 0 replicas after every migration. `stop_pods()` now marks every Deployment it stops with the + `wiremind.io/stopped-by` annotation; the restore only touches Deployments still carrying it. + ### Chore - Validate runtime/tooling compatibility on Python 3.14. +- CI: pushing a SemVer `vX.Y.Z-dev.N` tag (from any branch) publishes the PEP 440 dev + pre-release `X.Y.Z.devN` to PyPI, for validating changes on a canary environment before + merging. Stable releases still require a published GitHub Release. ## v6.0.0 (2025-10-07) diff --git a/VERSION b/VERSION index 024b066..66ce77b 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -6.2.1 +7.0.0 diff --git a/pyproject.toml b/pyproject.toml index 1ef768d..d181a8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" authors = [{ name = "wiremind", email = "dev@wiremind.io" }] license="LGPL-3.0-or-later" urls = { github = "https://github.com/wiremind/chartreuse"} -scripts = {chartreuse-upgrade = "chartreuse.chartreuse_upgrade:main"} +scripts = {chartreuse-upgrade = "chartreuse.chartreuse_upgrade:main", chartreuse-restore = "chartreuse.chartreuse_restore:main"} requires-python = ">=3.13.0" classifiers = [ "Programming Language :: Python :: 3", @@ -16,10 +16,10 @@ classifiers = [ dependencies = [ "alembic", + "kubernetes", "psycopg2", "pydantic>=2.0.0", "PyYAML", - "wiremind-kubernetes~=7.0", ] [build-system] diff --git a/requirements.txt b/requirements.txt index 84c02ef..77771ea 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,12 +16,10 @@ durationpy==0.10 # via kubernetes google-auth==2.41.1 # via kubernetes -greenlet==3.2.4 - # via sqlalchemy idna==3.10 # via requests kubernetes==34.1.0 - # via wiremind-kubernetes + # via chartreuse (pyproject.toml) mako==1.3.10 # via alembic markupsafe==3.0.3 @@ -75,5 +73,3 @@ urllib3==2.3.0 # requests websocket-client==1.9.0 # via kubernetes -wiremind-kubernetes==7.5.0 - # via chartreuse (pyproject.toml) diff --git a/src/chartreuse/chartreuse.py b/src/chartreuse/chartreuse.py index 7582986..f5f3611 100644 --- a/src/chartreuse/chartreuse.py +++ b/src/chartreuse/chartreuse.py @@ -1,8 +1,7 @@ import logging -import wiremind_kubernetes.kubernetes_helper - from .config_loader import DatabaseConfig +from .kubernetes_helper import KubernetesDeploymentManager from .utils import AlembicMigrationHelper logger = logging.getLogger(__name__) @@ -23,7 +22,7 @@ def __init__( self, databases_config: dict[str, DatabaseConfig], release_name: str, - kubernetes_helper: wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager | None = None, + kubernetes_helper: KubernetesDeploymentManager | None = None, ): configure_logging() @@ -56,9 +55,7 @@ def __init__( if kubernetes_helper: self.kubernetes_helper = kubernetes_helper else: - self.kubernetes_helper = wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager( - use_kubeconfig=None, release_name=release_name - ) + self.kubernetes_helper = KubernetesDeploymentManager(use_kubeconfig=None, release_name=release_name) @property def is_migration_needed(self) -> bool: diff --git a/src/chartreuse/chartreuse_restore.py b/src/chartreuse/chartreuse_restore.py new file mode 100644 index 0000000..d63caaa --- /dev/null +++ b/src/chartreuse/chartreuse_restore.py @@ -0,0 +1,36 @@ +import logging +import os + +from .chartreuse import configure_logging +from .chartreuse_upgrade import ensure_safe_run +from .kubernetes_helper import KubernetesDeploymentManager + +logger = logging.getLogger(__name__) + + +def main() -> None: + """ + Restore Deployments that a pre-deployment migration (CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT) + stopped and that the deployment itself could not scale back up. + + On that path chartreuse-upgrade intentionally skips start_pods() and relies on the + deployment that follows to restore replicas. That works for Deployments whose chart sets + spec.replicas, but HPA-managed ones omit it: they stay at 0 replicas and an HPA whose + target is at 0 with minReplicas >= 1 is ScalingDisabled and never scales them back up. + + Runs as a post-deployment (e.g. ArgoCD PostSync) hook: the new pod template is already + applied, so scaling up only ever starts new-version pods. Only Deployments still carrying + the annotation set by stop_pods() are touched. If this program fails, the deployment is + considered as failed: stranded workers must be visible, not a log line. + """ + configure_logging() + ensure_safe_run() + + release_name: str = os.environ["CHARTREUSE_RELEASE_NAME"] + + deployment_manager = KubernetesDeploymentManager(release_name=release_name, use_kubeconfig=None) + deployment_manager.restore_stopped_pods() + + +if __name__ == "__main__": + main() diff --git a/src/chartreuse/chartreuse_upgrade.py b/src/chartreuse/chartreuse_upgrade.py index 7c29079..81f5716 100755 --- a/src/chartreuse/chartreuse_upgrade.py +++ b/src/chartreuse/chartreuse_upgrade.py @@ -1,12 +1,11 @@ import logging import os -from wiremind_kubernetes import KubernetesDeploymentManager - from chartreuse import get_version from .chartreuse import Chartreuse from .config_loader import load_multi_database_config +from .kubernetes_helper import KubernetesDeploymentManager logger = logging.getLogger(__name__) diff --git a/src/chartreuse/kubernetes_helper.py b/src/chartreuse/kubernetes_helper.py new file mode 100644 index 0000000..558229e --- /dev/null +++ b/src/chartreuse/kubernetes_helper.py @@ -0,0 +1,375 @@ +""" +Kubernetes helpers to stop, start and restore the application's Deployments around a database +migration. + +Trimmed port of wiremind_kubernetes.KubernetesDeploymentManager (a library being +decommissioned): only what chartreuse uses survived — ExpectedDeploymentScale-driven +scale down/up with HPA disabling, plus the stopped-by annotation and restore logic. +Depends only on the official kubernetes client. +""" + +import functools +import logging +import os +import pprint +import time +from collections.abc import Callable, Generator +from typing import Any + +import kubernetes + +logger = logging.getLogger(__name__) + +HPA_ID_PREFIX = "wm--disabled--kube" +# Set by stop_pods() on every Deployment it scales down, cleared by start_pods() and +# restore_stopped_pods(). Allows telling "stopped by us" apart from "deliberately scaled to 0". +STOPPED_ANNOTATION = "wiremind.io/stopped-by" + + +def load_kubernetes_config(use_kubeconfig: bool | None = None) -> None: + """ + Load kubernetes configuration in memory, either from incluster method or from kubeconfig. + :param use_kubeconfig: + If True: Use ~/.kube/config file to authenticate. + If False: use kubernetes built-in incluster mechanism. + If None, will try to load built-in incluster mechanism, then try config file. + Defaults to None. + """ + if os.environ.get("CLASSIC_K8S_CONFIG"): + # We are in a Kind cluster for E2E test! Never use in cluster config. + kubernetes.config.load_kube_config() + return + + if use_kubeconfig is True: + kubernetes.config.load_kube_config() + elif use_kubeconfig is False: + kubernetes.config.load_incluster_config() + elif use_kubeconfig is None: + if os.path.exists(kubernetes.config.incluster_config.SERVICE_TOKEN_FILENAME): + kubernetes.config.load_incluster_config() + else: + kubernetes.config.load_kube_config() + logger.debug("Kubernetes configuration successfully set.") + + +def retry_kubernetes_request(function: Callable) -> Callable: + """ + Decorator that retries a failed Kubernetes API request if needed and ignores 404 + """ + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return function(*args, **kwargs) + except kubernetes.client.rest.ApiException as e: + if e.status == 404: + logger.warning("Not found, ignoring.") + return + logger.error(e) + logger.info("Retrying in 5 seconds...") + time.sleep(5) + return function(*args, **kwargs) + + return wrapper + + +def retry_kubernetes_request_no_ignore(function: Callable) -> Callable: + """ + Decorator that retries a failed Kubernetes API request if needed and do NOT ignore 404 (raise if 404) + """ + + @functools.wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return function(*args, **kwargs) + except kubernetes.client.rest.ApiException as e: + if e.status == 404: + raise + logger.error(e) + logger.info("Retrying in 5 seconds...") + time.sleep(5) + return function(*args, **kwargs) + + return wrapper + + +def _get_namespace_from_kube() -> str: + return open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read() + + +class KubernetesDeploymentManager: + """ + Scale down/up all Deployments that should be stopped/started when doing database + migration/maintenance (alembic, dump, etc). Each managed Deployment must define an + ExpectedDeploymentScale. + + Usage: + manager = KubernetesDeploymentManager(release_name="my-release", use_kubeconfig=None) + manager.stop_pods() + do_something('wololo') + manager.start_pods() + """ + + SCALE_DOWN_MAX_WAIT_TIME: int = 3600 + + def __init__( + self, + release_name: str, + use_kubeconfig: bool | None = False, + namespace: str | None = None, + should_load_kubernetes_config: bool = True, + ): + if should_load_kubernetes_config: + load_kubernetes_config(use_kubeconfig=use_kubeconfig) + self.release_name = release_name + self.namespace = namespace if namespace else _get_namespace_from_kube() + self.client_appsv1_api = kubernetes.client.AppsV1Api() + self.client_corev1_api = kubernetes.client.CoreV1Api() + self.client_autoscalingv2_api = kubernetes.client.AutoscalingV2Api() + self.client_custom_objects_api = kubernetes.client.CustomObjectsApi() + + def get_deployment_scale(self, deployment_name: str) -> kubernetes.client.V1Scale: + logger.debug("Getting deployment scale for %s", deployment_name) + return self.client_appsv1_api.read_namespaced_deployment_scale(deployment_name, self.namespace) + + @retry_kubernetes_request + def scale_down_deployment(self, deployment_name: str) -> None: + body = self.get_deployment_scale(deployment_name) + logger.debug("Deleting all Pods for %s", deployment_name) + body.spec.replicas = 0 + self.client_appsv1_api.patch_namespaced_deployment_scale(deployment_name, self.namespace, body) + logger.debug("Done deleting.") + + @retry_kubernetes_request + def scale_up_deployment(self, deployment_name: str, pod_amount: int) -> None: + body = self.get_deployment_scale(deployment_name) + logger.debug("Recreating backend Pods for %s", deployment_name) + body.spec.replicas = pod_amount + self.client_appsv1_api.patch_namespaced_deployment_scale(deployment_name, self.namespace, body) + logger.debug("Done recreating.") + + @retry_kubernetes_request_no_ignore + def _get_pods_from_deployment(self, deployment_name: str) -> list: + logger.debug("Asking if Deployment %s is stopped", deployment_name) + labels = self.client_appsv1_api.read_namespaced_deployment( + deployment_name, self.namespace + ).spec.selector.match_labels + + try: + return self.client_corev1_api.list_namespaced_pod( + namespace=self.namespace, + label_selector=",".join(["{}={}".format(*kv) for kv in labels.items()]), + ).items + except kubernetes.client.rest.ApiException as e: + if e.status == 404: + return [] + else: + raise + + def is_deployment_stopped(self, deployment_name: str) -> bool: + try: + pod_list: list = self._get_pods_from_deployment(deployment_name) + except kubernetes.client.rest.ApiException as e: + if e.status == 404: + logger.warning("Not found, ignoring.") + return True + raise + + current_scale = 0 + for pod in pod_list: + if pod.status.phase not in ("Failed"): + current_scale += 1 + + if current_scale > 0: + logger.info("%s Deployment has %s living replicas", deployment_name, current_scale) + return False + return True + + def get_deployment_hpa(self, *, deployment_name: str) -> Generator: + for hpa in self.client_autoscalingv2_api.list_namespaced_horizontal_pod_autoscaler(self.namespace).items: + if hpa.spec.scale_target_ref.kind == "Deployment" and hpa.spec.scale_target_ref.name == deployment_name: + yield hpa + + def patch_deployment_hpa(self, *, hpa_name: str, body: Any) -> None: + self.client_autoscalingv2_api.patch_namespaced_horizontal_pod_autoscaler( + name=hpa_name, namespace=self.namespace, body=body + ) + + @retry_kubernetes_request + def annotate_deployment(self, deployment_name: str, annotations: dict[str, str | None]) -> None: + """ + Add or update annotations on a Deployment. A None value removes the annotation. + """ + body = {"metadata": {"annotations": annotations}} + self.client_appsv1_api.patch_namespaced_deployment(deployment_name, self.namespace, body) + + @retry_kubernetes_request_no_ignore + def _get_expected_deployment_scale_dict(self) -> dict[int, dict[str, int]]: + """ + Return a dict of expected deployment scale: + { + 0: { # priority + # key: Deployment name, only if it has an associated eds + # value: expected Deployment Scale (replicas) + "my-deployment": 3, + "my-other-deployment": 42 + }, + 1: { + "my-third-deployment": 17, + }, + } + """ + logger.debug("Getting Expected Deployment Scale list") + eds_list: list[dict[str, Any]] = [] + release_label_keys = ["app.kubernetes.io/instance", "release"] + + for release_label_key in release_label_keys: + logger.debug(f"Getting Expected Deployment Scale list with the release label key {release_label_key}") + try: + eds_list.extend( + self.client_custom_objects_api.list_namespaced_custom_object( + namespace=self.namespace, + group="wiremind.io", + version="v1", + plural="expecteddeploymentscales", + label_selector=f"{release_label_key}={self.release_name}", + )["items"] + ) + except kubernetes.client.rest.ApiException as e: + if e.status != 404: + raise + + eds_dict: dict[int, dict[str, int]] = {} + for eds in eds_list: + deployment_name: str = eds["spec"]["deploymentName"] + expected_scale: int = eds["spec"]["expectedScale"] + priority: int = eds["spec"].get("priority", 0) + + if priority not in eds_dict: + eds_dict[priority] = {} + + eds_dict[priority][deployment_name] = expected_scale + + logger.debug("Deployments are %s", pprint.pformat(eds_dict)) + return eds_dict + + @retry_kubernetes_request + def disable_hpa(self, *, deployment_name: str) -> None: + for hpa in self.get_deployment_hpa(deployment_name=deployment_name): + # Tell the hpa to manage a non-existing Deployment + hpa.spec.scale_target_ref.name = f"{HPA_ID_PREFIX}-{deployment_name}" + self.patch_deployment_hpa(hpa_name=hpa.metadata.name, body=hpa) + + @retry_kubernetes_request + def re_enable_hpa(self, *, deployment_name: str) -> None: + for hpa in self.get_deployment_hpa(deployment_name=f"{HPA_ID_PREFIX}-{deployment_name}"): + hpa.spec.scale_target_ref.name = deployment_name + self.patch_deployment_hpa(hpa_name=hpa.metadata.name, body=hpa) + + def _are_deployments_stopped(self, deployment_dict: dict[str, int]) -> bool: + for deployment_name in deployment_dict: + if not self.is_deployment_stopped(deployment_name): + return False + return True + + def _stop_deployments(self, deployment_dict: dict[str, int]) -> None: + """ + Scale down a dict (deployment_name, expected_scale) of Deployments. + """ + for deployment_name in deployment_dict: + self.annotate_deployment(deployment_name, {STOPPED_ANNOTATION: self.release_name}) + for _ in range(self.SCALE_DOWN_MAX_WAIT_TIME): + for deployment_name in deployment_dict: + self.disable_hpa(deployment_name=deployment_name) + self.scale_down_deployment(deployment_name) + if self._are_deployments_stopped(deployment_dict): + break + time.sleep(1) + else: + raise Exception("Timed out waiting for pods to be deleted: aborting.") + + def stop_pods(self) -> None: + """ + Scale to 0 all deployments for which an ExpectedDeploymentScale links to. + stop all deployments, then wait for actual stop, by priority (descending order): + Example: stop all deployments with priority 1, then all deployments with priority 0 + """ + expected_deployment_scale_dict: dict[int, dict[str, int]] = self._get_expected_deployment_scale_dict() + + logger.info("Scaling down application Deployments...") + if not expected_deployment_scale_dict: + logger.info("No Deployments to scale down") + return + + priorities: list[int] = sorted(expected_deployment_scale_dict, reverse=True) + for priority in priorities: + priority_dict: dict[str, int] = expected_deployment_scale_dict[priority] + if len(priority_dict): + self._stop_deployments(priority_dict) + logger.info("Done scaling down application Deployments.") + + def start_pods(self) -> None: + """ + Start all Pods that should be started + """ + expected_deployment_scale_dict: dict[int, dict[str, int]] = self._get_expected_deployment_scale_dict() + + logger.info("Scaling up application Deployments...") + if not expected_deployment_scale_dict: + logger.info("No Deployments to scale up") + return + + scaled: bool = False + for priority_dict in expected_deployment_scale_dict.values(): + if len(priority_dict): + scaled = True + for name, expected_scale in priority_dict.items(): + self.re_enable_hpa(deployment_name=name) + self.scale_up_deployment(name, expected_scale) + self.annotate_deployment(name, {STOPPED_ANNOTATION: None}) + if scaled: + logger.info("Done scaling up application Deployments") + else: + logger.info("No Deployments to scale up") + + def restore_stopped_pods(self) -> None: + """ + Restore Deployments that stop_pods() scaled down and that nothing scaled back up since. + + Only Deployments still carrying STOPPED_ANNOTATION are touched, so a Deployment + deliberately scaled to zero outside of stop_pods() is left alone. + + Meant to run after the deployment that follows a stop_pods() without start_pods() + (CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT): that deployment restores spec.replicas of + Deployments that define it in their chart, but HPA-managed ones omit it, and an HPA + whose target is at 0 replicas with minReplicas >= 1 is ScalingDisabled and will never + scale it back up. + """ + expected_deployment_scale_dict: dict[int, dict[str, int]] = self._get_expected_deployment_scale_dict() + + logger.info("Restoring Deployments left stopped...") + restored: list[str] = [] + for priority in sorted(expected_deployment_scale_dict): + for name, expected_scale in expected_deployment_scale_dict[priority].items(): + try: + deployment = self.client_appsv1_api.read_namespaced_deployment(name, self.namespace) + except kubernetes.client.rest.ApiException as e: + if e.status == 404: + continue + raise + if STOPPED_ANNOTATION not in (deployment.metadata.annotations or {}): + continue + self.re_enable_hpa(deployment_name=name) + if not deployment.spec.replicas: + target_scale = expected_scale + if target_scale < 1 and any(self.get_deployment_hpa(deployment_name=name)): + # Any scale >= 1 re-arms the HPA, which then enforces its own minReplicas + target_scale = 1 + if target_scale > 0: + self.scale_up_deployment(name, target_scale) + self.annotate_deployment(name, {STOPPED_ANNOTATION: None}) + restored.append(name) + if restored: + logger.info("Done restoring stopped Deployments: %s", ", ".join(restored)) + else: + logger.info("No stopped Deployments to restore") diff --git a/src/chartreuse/tests/conftest.py b/src/chartreuse/tests/conftest.py index 1232161..9085588 100644 --- a/src/chartreuse/tests/conftest.py +++ b/src/chartreuse/tests/conftest.py @@ -59,9 +59,9 @@ def configure_chartreuse_mock( # Mock KubernetesDeploymentManager in all the right places mock_kdm = mocker.MagicMock() - mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager", return_value=mock_kdm) + mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager", return_value=mock_kdm) mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager", return_value=mock_kdm) # Also mock the kubernetes config loading to prevent the config errors mocker.patch("kubernetes.config.load_incluster_config") - mocker.patch("wiremind_kubernetes.kube_config.load_kubernetes_config") + mocker.patch("chartreuse.kubernetes_helper.load_kubernetes_config") diff --git a/src/chartreuse/tests/e2e_tests/conftest.py b/src/chartreuse/tests/e2e_tests/conftest.py index bed5fe4..625eeca 100644 --- a/src/chartreuse/tests/e2e_tests/conftest.py +++ b/src/chartreuse/tests/e2e_tests/conftest.py @@ -1,22 +1,63 @@ import logging import os import subprocess +import sys import time +import urllib.parse from collections.abc import Generator import pytest import sqlalchemy from sqlalchemy import inspect -from wiremind_kubernetes.kube_config import load_kubernetes_config -from wiremind_kubernetes.kubernetes_helper import KubernetesDeploymentManager -from wiremind_kubernetes.tests.e2e_tests.conftest import create_namespace, setUpE2E # noqa: F401 -from wiremind_kubernetes.utils import run_command import chartreuse +from chartreuse.kubernetes_helper import KubernetesDeploymentManager, load_kubernetes_config +from chartreuse.utils.command import run_command TEST_NAMESPACE = "chartreuse-e2e-test" TEST_RELEASE = "e2e-test-release" +# Test clusters the e2e suite is allowed to run against +TEST_IPS_WHITELISTED = [ + "localhost", # minikube + "127.0.0.1", # kind + "kubernetes.docker.internal", # Docker for Mac +] +TEST_NODES_WHITELISTED = [ + "minikube", + "kind-control-plane", + "kind", + "kind-worker", +] + + +def _is_ip_whitelisted() -> bool: + api_server = subprocess.check_output( + "kubectl config view --minify | grep server | cut -f 2- -d ':' | tr -d ' '", + shell=True, + text=True, + ) + hostname = urllib.parse.urlparse(api_server.lower().strip()).hostname + return hostname in TEST_IPS_WHITELISTED + + +def _is_node_whitelisted() -> bool: + output, *_ = run_command("kubectl get nodes -o name", return_result=True) + cluster_nodes = [x for x in output.replace("node/", "").split("\n") if x != ""] + return all(node in TEST_NODES_WHITELISTED for node in cluster_nodes) + + +@pytest.fixture(scope="session", autouse=True) +def setUpE2E() -> None: + """ + sys.exit(1) if kubectl current context api server is not a test cluster (like kind, minikube, etc) + """ + logging.info("[CLUSTER-CONFIG]: Making sure the tests are running against a test cluster...") + if not _is_ip_whitelisted() and not _is_node_whitelisted(): + logging.error("Attempted to run tests with a non-test cluster, aborting!") + sys.exit(1) + + ROOT_PATH = os.path.join(os.path.dirname(chartreuse.__file__), "..", "..") EXAMPLE_PATH = os.path.join(ROOT_PATH, "example") HELM_CHART_PATH = os.path.join(EXAMPLE_PATH, "helm-chart", "my-example-chart") diff --git a/src/chartreuse/tests/unit_tests/conftest.py b/src/chartreuse/tests/unit_tests/conftest.py index 5562b10..4ba5b2a 100644 --- a/src/chartreuse/tests/unit_tests/conftest.py +++ b/src/chartreuse/tests/unit_tests/conftest.py @@ -12,6 +12,6 @@ def configure_chartreuse_mock(mocker: MockerFixture, is_migration_needed: bool = ) mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse.upgrade") mocker.patch( - "wiremind_kubernetes.kubernetes_helper._get_namespace_from_kube", + "chartreuse.kubernetes_helper._get_namespace_from_kube", return_value="foo", ) diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse.py b/src/chartreuse/tests/unit_tests/test_chartreuse.py index d30e074..7c35b2a 100644 --- a/src/chartreuse/tests/unit_tests/test_chartreuse.py +++ b/src/chartreuse/tests/unit_tests/test_chartreuse.py @@ -87,7 +87,7 @@ def test_chartreuse_init_without_kubernetes_helper(self, mocker: MockerFixture) mock_alembic_helper.return_value = mock_alembic_instance # Mock KubernetesDeploymentManager - mock_k8s_manager = mocker.patch("wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager") + mock_k8s_manager = mocker.patch("chartreuse.chartreuse.KubernetesDeploymentManager") mock_k8s_instance = MagicMock() mock_k8s_manager.return_value = mock_k8s_instance @@ -315,7 +315,7 @@ def test_multi_chartreuse_init_without_kubernetes_helper(self, mocker: MockerFix mock_alembic_helper.return_value = mock_alembic_instance # Mock KubernetesDeploymentManager - mock_k8s_manager = mocker.patch("wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager") + mock_k8s_manager = mocker.patch("chartreuse.chartreuse.KubernetesDeploymentManager") mock_k8s_instance = MagicMock() mock_k8s_manager.return_value = mock_k8s_instance diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse_restore.py b/src/chartreuse/tests/unit_tests/test_chartreuse_restore.py new file mode 100644 index 0000000..cae389f --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_chartreuse_restore.py @@ -0,0 +1,58 @@ +import pytest +from pytest_mock.plugin import MockerFixture + +import chartreuse.chartreuse_restore + +from ..conftest import configure_os_environ_mock + + +def _configure_deployment_manager_mock(mocker: MockerFixture) -> "MockerFixture": + mocked_kdm = mocker.MagicMock() + mocker.patch("chartreuse.chartreuse_restore.KubernetesDeploymentManager", return_value=mocked_kdm) + mocker.patch( + "chartreuse.kubernetes_helper._get_namespace_from_kube", + return_value="foo", + ) + return mocked_kdm + + +def test_chartreuse_restore_restores_stopped_pods(mocker: MockerFixture) -> None: + """ + Test that chartreuse_restore restores the Deployments that stop_pods left at 0 replicas. + """ + mocked_kdm = _configure_deployment_manager_mock(mocker) + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") + configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_APP_VERSION": "5.0.0"}) + + chartreuse.chartreuse_restore.main() + + mocked_kdm.restore_stopped_pods.assert_called_once() + + +def test_chartreuse_restore_incompatible_helm_chart_version(mocker: MockerFixture) -> None: + """ + Test that chartreuse_restore aborts before touching anything when the package and the + Helm Chart don't have the same major.minor. + """ + mocked_kdm = _configure_deployment_manager_mock(mocker) + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.1.0") + configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_APP_VERSION": "5.0.0"}) + + with pytest.raises(ValueError): + chartreuse.chartreuse_restore.main() + + mocked_kdm.restore_stopped_pods.assert_not_called() + + +def test_chartreuse_restore_failure_is_fatal(mocker: MockerFixture) -> None: + """ + Test that a restore failure propagates: the Job must fail so the deployment is marked + as failed, instead of leaving workers silently stranded at 0 replicas. + """ + mocked_kdm = _configure_deployment_manager_mock(mocker) + mocked_kdm.restore_stopped_pods.side_effect = Exception("boom") + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") + configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_APP_VERSION": "5.0.0"}) + + with pytest.raises(Exception, match="boom"): + chartreuse.chartreuse_restore.main() diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py index a17d7da..d3df7d7 100644 --- a/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py +++ b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py @@ -14,8 +14,8 @@ def test_chartreuse_upgrade_detected_migration_enabled_stop_pods( Test that chartreuse_upgrades stop pods in case of detected migration. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=True) - mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") - mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") + mocked_stop_pods = mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.stop_pods") + mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.start_pods") mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_APP_VERSION": "5.0.0"}) @@ -30,8 +30,8 @@ def test_chartreuse_upgrade_detected_migration_disabled_stop_pods( Test that chartreuse_upgrades does not stop pods in case of detected migration but we disallow stop-pods. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=True) - mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") - mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") + mocked_stop_pods = mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.stop_pods") + mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.start_pods") mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") configure_os_environ_mock( mocker=mocker, @@ -49,8 +49,8 @@ def test_chartreuse_upgrade_no_migration_disabled_stop_pods( Test that chartreuse_upgrades does NOT stop pods in case of migration not needed. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=False) - mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") - mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") + mocked_stop_pods = mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.stop_pods") + mocker.patch("chartreuse.kubernetes_helper.KubernetesDeploymentManager.start_pods") mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_APP_VERSION": "5.0.0"}) diff --git a/src/chartreuse/tests/unit_tests/test_kubernetes_helper.py b/src/chartreuse/tests/unit_tests/test_kubernetes_helper.py new file mode 100644 index 0000000..053eca2 --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_kubernetes_helper.py @@ -0,0 +1,173 @@ +import unittest + +import kubernetes +from pytest_mock import MockerFixture + +from chartreuse.kubernetes_helper import STOPPED_ANNOTATION, KubernetesDeploymentManager + + +def _make_kdm(mocker: MockerFixture) -> KubernetesDeploymentManager: + mocker.patch("kubernetes.client.AppsV1Api") + mocker.patch("kubernetes.client.CoreV1Api") + mocker.patch("kubernetes.client.AutoscalingV2Api") + mocker.patch("kubernetes.client.CustomObjectsApi") + return KubernetesDeploymentManager( + should_load_kubernetes_config=False, + namespace="foo", + release_name="concerned", + ) + + +def _make_deployment(mocker: MockerFixture, annotations: dict | None, replicas: int) -> unittest.mock.MagicMock: + deployment = mocker.MagicMock() + deployment.metadata.annotations = annotations + deployment.spec.replicas = replicas + return deployment + + +def test_stop_pods_priority(mocker: MockerFixture) -> None: + """ + Test that we honor priorities when stopping workloads. + + i.e that we stop all deployments with the highest priority first, then wait for all of those to be stopped + then continue. + """ + kdm = _make_kdm(mocker) + mocker.patch.object( + kdm, + "_get_expected_deployment_scale_dict", + return_value={ + 0: {"last": 42}, + 2: {"first": 17}, + 1: {"second": 17}, + }, + ) + mocked_stop_deployments = mocker.patch.object(kdm, "_stop_deployments") + + kdm.stop_pods() + + assert mocked_stop_deployments.mock_calls == [ + unittest.mock.call({"first": 17}), + unittest.mock.call({"second": 17}), + unittest.mock.call({"last": 42}), + ] + + +def test_stop_deployments_correctly_wait(mocker: MockerFixture) -> None: + """ + Test that we wait for deployments to be stopped + """ + kdm = _make_kdm(mocker) + deployment_dict = {"my-pod": 42, "my-other-pod": 113} + mocker.patch.object(kdm, "disable_hpa") + mocker.patch.object(kdm, "scale_down_deployment") + mocker.patch.object(kdm, "annotate_deployment") + mocked_are_deployments_stopped = mocker.patch.object( + kdm, "_are_deployments_stopped", side_effect=[False, False, True] + ) + mocker.patch("time.sleep") + + kdm._stop_deployments(deployment_dict) + + assert mocked_are_deployments_stopped.mock_calls == [ + unittest.mock.call(deployment_dict), + unittest.mock.call(deployment_dict), + unittest.mock.call(deployment_dict), + ] + + +def test_stop_deployments_annotates_stopped_deployments(mocker: MockerFixture) -> None: + """ + Test that _stop_deployments marks every Deployment it scales down with STOPPED_ANNOTATION. + """ + kdm = _make_kdm(mocker) + mocker.patch.object(kdm, "_are_deployments_stopped", return_value=True) + mocker.patch.object(kdm, "disable_hpa") + mocker.patch.object(kdm, "scale_down_deployment") + mocked_annotate = mocker.patch.object(kdm, "annotate_deployment") + + kdm._stop_deployments({"my-pod": 42, "my-other-pod": 113}) + + assert mocked_annotate.mock_calls == [ + unittest.mock.call("my-pod", {STOPPED_ANNOTATION: "concerned"}), + unittest.mock.call("my-other-pod", {STOPPED_ANNOTATION: "concerned"}), + ] + + +def test_start_pods_clears_stopped_annotation(mocker: MockerFixture) -> None: + kdm = _make_kdm(mocker) + mocker.patch.object(kdm, "_get_expected_deployment_scale_dict", return_value={0: {"worker": 2}}) + mocker.patch.object(kdm, "re_enable_hpa") + mocker.patch.object(kdm, "scale_up_deployment") + mocked_annotate = mocker.patch.object(kdm, "annotate_deployment") + + kdm.start_pods() + + assert mocked_annotate.mock_calls == [unittest.mock.call("worker", {STOPPED_ANNOTATION: None})] + + +def test_restore_stopped_pods(mocker: MockerFixture) -> None: + """ + Test that restore_stopped_pods only restores Deployments still carrying STOPPED_ANNOTATION, + only scales those still at 0 replicas, and always clears the annotation. + """ + kdm = _make_kdm(mocker) + mocker.patch.object( + kdm, + "_get_expected_deployment_scale_dict", + return_value={ + 0: { + "stopped-worker": 2, # annotated, still at 0: restore to expectedScale + "untouched-worker": 3, # not annotated (deliberately at 0): leave alone + "already-scaled": 1, # annotated but replicas > 0: only clear the annotation + }, + 1: {"gone": 1}, # Deployment does not exist anymore: skip + }, + ) + deployments = { + "stopped-worker": _make_deployment(mocker, {STOPPED_ANNOTATION: "concerned"}, 0), + "untouched-worker": _make_deployment(mocker, None, 0), + "already-scaled": _make_deployment(mocker, {STOPPED_ANNOTATION: "concerned"}, 4), + } + + def read_deployment(name: str, namespace: str) -> unittest.mock.MagicMock: + if name not in deployments: + raise kubernetes.client.rest.ApiException(status=404) + return deployments[name] + + kdm.client_appsv1_api.read_namespaced_deployment.side_effect = read_deployment + mocked_re_enable_hpa = mocker.patch.object(kdm, "re_enable_hpa") + mocked_scale_up = mocker.patch.object(kdm, "scale_up_deployment") + mocked_annotate = mocker.patch.object(kdm, "annotate_deployment") + + kdm.restore_stopped_pods() + + assert mocked_re_enable_hpa.mock_calls == [ + unittest.mock.call(deployment_name="stopped-worker"), + unittest.mock.call(deployment_name="already-scaled"), + ] + assert mocked_scale_up.mock_calls == [unittest.mock.call("stopped-worker", 2)] + assert mocked_annotate.mock_calls == [ + unittest.mock.call("stopped-worker", {STOPPED_ANNOTATION: None}), + unittest.mock.call("already-scaled", {STOPPED_ANNOTATION: None}), + ] + + +def test_restore_stopped_pods_rearms_hpa_when_expected_scale_is_zero(mocker: MockerFixture) -> None: + """ + An EDS expectedScale of 0 cannot re-arm an HPA with minReplicas >= 1: restore to 1 instead + and let the HPA enforce its own minimum. + """ + kdm = _make_kdm(mocker) + mocker.patch.object(kdm, "_get_expected_deployment_scale_dict", return_value={0: {"worker": 0}}) + kdm.client_appsv1_api.read_namespaced_deployment.return_value = _make_deployment( + mocker, {STOPPED_ANNOTATION: "concerned"}, 0 + ) + mocker.patch.object(kdm, "re_enable_hpa") + mocker.patch.object(kdm, "get_deployment_hpa", return_value=iter([mocker.MagicMock()])) + mocked_scale_up = mocker.patch.object(kdm, "scale_up_deployment") + mocker.patch.object(kdm, "annotate_deployment") + + kdm.restore_stopped_pods() + + assert mocked_scale_up.mock_calls == [unittest.mock.call("worker", 1)] diff --git a/src/chartreuse/utils/alembic_migration_helper.py b/src/chartreuse/utils/alembic_migration_helper.py index efc7d4c..57adf3c 100755 --- a/src/chartreuse/utils/alembic_migration_helper.py +++ b/src/chartreuse/utils/alembic_migration_helper.py @@ -5,7 +5,8 @@ import sqlalchemy from sqlalchemy import inspect -from wiremind_kubernetes.utils import run_command + +from .command import run_command logger = logging.getLogger(__name__) diff --git a/src/chartreuse/utils/command.py b/src/chartreuse/utils/command.py new file mode 100644 index 0000000..6bab0db --- /dev/null +++ b/src/chartreuse/utils/command.py @@ -0,0 +1,40 @@ +import logging +import shlex +import subprocess + +logger = logging.getLogger(__name__) + + +def run_command(command: list | str, return_result: bool = False, **kw_args) -> tuple[str, str, int]: + """ + Run command, print stdout/stderr, check that command exited correctly, return stdout/err + """ + logger.debug(f"Running {command}") + + interpreted_command: list[str] + if isinstance(command, str): + interpreted_command = shlex.split(command) + else: + interpreted_command = command + + process = subprocess.Popen( + interpreted_command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + **kw_args, + ) + + if return_result: + out, err = process.communicate() + return (out, err, process.returncode) + + if process.stdout: + for line in iter(process.stdout.readline, ""): + logger.info(line.strip()) + process.wait() + + if process.returncode: + raise subprocess.CalledProcessError(process.returncode, command) + + return "", "", 0