Add statefulset_is_fully_available and statefulset_is_not_fully_available

Rodrigo Menezes · Lawouach · commit e05be6094ec2 · 2020-09-23T12:05:23.000+02:00
Signed-off-by: Rodrigo Menezes &lt;rodrigomenezes@pinterest.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,11 +2,15 @@
 
 ## [Unreleased][]
 
+[Unreleased]: https://github.com/chaostoolkit/chaostoolkit-kubernetes/compare/0.22.0...HEAD
+
+### Added
+
 - Add namespace to all probes and actions with backward compatibility [#39][39]  
+- Added `statefulset_is_fully_available` and `statefulset_is_not_fully_available` [#86][86].
 
 [39]: https://github.com/chaostoolkit/chaostoolkit-kubernetes/issues/39
-
-[Unreleased]: https://github.com/chaostoolkit/chaostoolkit-kubernetes/compare/0.22.0...HEAD
+[86]: https://github.com/chaostoolkit/chaostoolkit-kubernetes/pull/86
 
 ### Changed
 
diff --git a/chaosk8s/__init__.py b/chaosk8s/__init__.py
@@ -142,6 +142,7 @@ def load_exported_activities() -> List[DiscoveredActivities]:
     activities.extend(discover_actions("chaosk8s.service.actions"))
     activities.extend(discover_actions("chaosk8s.service.probes"))
     activities.extend(discover_actions("chaosk8s.statefulset.actions"))
+    activities.extend(discover_probes("chaosk8s.statefulset.probes"))
     return activities
 
 
diff --git a/chaosk8s/statefulset/probes.py b/chaosk8s/statefulset/probes.py
@@ -0,0 +1,123 @@
+# -*- coding: utf-8 -*-
+from typing import Union
+
+import urllib3
+from chaoslib.exceptions import ActivityFailed
+from chaoslib.types import Secrets
+from functools import partial
+from kubernetes import client, watch
+from logzero import logger
+
+from chaosk8s import create_k8s_api_client
+
+__all__ = [
+    "statefulset_fully_available",
+    "statefulset_not_fully_available"
+]
+
+
+def _statefulset_readiness_has_state(name: str, ready: bool,
+                                     ns: str = "default",
+                                     label_selector: str = None,
+                                     timeout: int = 30,
+                                     secrets: Secrets = None):
+    """
+    Check wether if the given statefulSet state is ready or not
+    according to the ready paramter.
+    If the state is not reached after `timeout` seconds, a
+    :exc:`chaoslib.exceptions.ActivityFailed` exception is raised.
+    """
+    field_selector = "metadata.name={name}".format(name=name)
+    api = create_k8s_api_client(secrets)
+    v1 = client.AppsV1Api(api)
+    w = watch.Watch()
+    timeout = int(timeout)
+
+    if label_selector is None:
+        watch_events = partial(w.stream, v1.list_namespaced_stateful_set,
+                               namespace=ns,
+                               field_selector=field_selector,
+                               _request_timeout=timeout)
+    else:
+        label_selector = label_selector.format(name=name)
+        watch_events = partial(w.stream, v1.list_namespaced_stateful_set,
+                               namespace=ns,
+                               field_selector=field_selector,
+                               label_selector=label_selector,
+                               _request_timeout=timeout)
+
+    try:
+        logger.debug("Watching events for {t}s".format(t=timeout))
+        for event in watch_events():
+            statefulset = event['object']
+            status = statefulset.status
+            spec = statefulset.spec
+
+            logger.debug(
+                "StatefulSet '{p}' {t}: "
+                "Ready Replicas {r} - "
+                "Unavailable Replicas {u} - "
+                "Desired Replicas {a}".format(
+                    p=statefulset.metadata.name, t=event["type"],
+                    r=status.ready_replicas,
+                    a=spec.replicas,
+                    u=status.unavailable_replicas))
+
+            readiness = status.ready_replicas == spec.replicas
+            if ready == readiness:
+                w.stop()
+                return True
+
+    except urllib3.exceptions.ReadTimeoutError:
+        logger.debug("Timed out!")
+        return False
+
+
+def statefulset_not_fully_available(name: str, ns: str = "default",
+                                    label_selector: str = None,
+                                    timeout: int = 30,
+                                    secrets: Secrets = None):
+    """
+    Wait until the statefulSet gets into an intermediate state where not all
+    expected replicas are available. Once this state is reached, return `True`.
+    If the state is not reached after `timeout` seconds, a
+    :exc:`chaoslib.exceptions.ActivityFailed` exception is raised.
+    """
+    if _statefulset_readiness_has_state(
+        name,
+        False,
+        ns,
+        label_selector,
+        timeout,
+        secrets,
+    ):
+        return True
+    else:
+        raise ActivityFailed(
+            "microservice '{name}' failed to stop running within {t}s".format(
+                name=name, t=timeout))
+
+
+def statefulset_fully_available(name: str, ns: str = "default",
+                                label_selector: str = None,
+                                timeout: int = 30,
+                                secrets: Secrets = None):
+    """
+    Wait until all the statefulSet expected replicas are available.
+    Once this state is reached, return `True`.
+    If the state is not reached after `timeout` seconds, a
+    :exc:`chaoslib.exceptions.ActivityFailed` exception is raised.
+    """
+    if _statefulset_readiness_has_state(
+        name,
+        True,
+        ns,
+        label_selector,
+        timeout,
+        secrets,
+    ):
+        return True
+    else:
+        raise ActivityFailed(
+            "microservice '{name}' failed to recover within {t}s".format(
+                name=name, t=timeout))
diff --git a/tests/test_probes.py b/tests/test_probes.py
@@ -12,6 +12,8 @@
     microservice_available_and_healthy, microservice_is_not_available, \
     service_endpoint_is_initialized, deployment_is_not_fully_available, \
     deployment_is_fully_available, read_microservices_logs
+from chaosk8s.statefulset.probes import statefulset_fully_available, \
+    statefulset_not_fully_available
 
 
 @patch('chaosk8s.has_local_config_file', autospec=True)
@@ -223,6 +225,88 @@ def test_deployment_is_not_fully_available_when_it_should(cl, client,
         str(excinfo.value)
 
 
+@patch('chaosk8s.has_local_config_file', autospec=True)
+@patch('chaosk8s.statefulset.probes.watch', autospec=True)
+@patch('chaosk8s.statefulset.probes.client', autospec=True)
+@patch('chaosk8s.client')
+def test_statefulset_not_fully_available(cl, client, watch, has_conf):
+    has_conf.return_value = False
+    statefulset = MagicMock()
+    statefulset.spec.replicas = 2
+    statefulset.status.ready_replicas = 1
+
+    watcher = MagicMock()
+    watcher.stream = MagicMock()
+    watcher.stream.side_effect = [[{"object": statefulset, "type": "ADDED"}]]
+    watch.Watch.return_value = watcher
+
+    assert statefulset_not_fully_available("mysvc") is True
+
+
+@patch('chaosk8s.has_local_config_file', autospec=True)
+@patch('chaosk8s.statefulset.probes.watch', autospec=True)
+@patch('chaosk8s.statefulset.probes.client', autospec=True)
+@patch('chaosk8s.client')
+def test_statefulset_fully_available_when_it_should_not(cl, client,
+                                                          watch, has_conf):
+    has_conf.return_value = False
+    statefulset = MagicMock()
+    statefulset.spec.replicas = 2
+    statefulset.status.ready_replicas = 2
+
+    watcher = MagicMock()
+    watcher.stream = MagicMock()
+    watcher.stream.side_effect = urllib3.exceptions.ReadTimeoutError(
+        None, None, None)
+    watch.Watch.return_value = watcher
+
+    with pytest.raises(ActivityFailed) as excinfo:
+        statefulset_not_fully_available("mysvc")
+    assert "microservice 'mysvc' failed to stop running within" in \
+        str(excinfo.value)
+
+
+@patch('chaosk8s.has_local_config_file', autospec=True)
+@patch('chaosk8s.statefulset.probes.watch', autospec=True)
+@patch('chaosk8s.statefulset.probes.client', autospec=True)
+@patch('chaosk8s.client')
+def test_statefulset_fully_available(cl, client, watch, has_conf):
+    has_conf.return_value = False
+    statefulset = MagicMock()
+    statefulset.spec.replicas = 2
+    statefulset.status.ready_replicas = 2
+
+    watcher = MagicMock()
+    watcher.stream = MagicMock()
+    watcher.stream.side_effect = [[{"object": statefulset, "type": "ADDED"}]]
+    watch.Watch.return_value = watcher
+
+    assert statefulset_fully_available("mysvc") is True
+
+
+@patch('chaosk8s.has_local_config_file', autospec=True)
+@patch('chaosk8s.statefulset.probes.watch', autospec=True)
+@patch('chaosk8s.statefulset.probes.client', autospec=True)
+@patch('chaosk8s.client')
+def test_statefulset_not_fully_available_when_it_should(cl, client,
+                                                          watch, has_conf):
+    has_conf.return_value = False
+    statefulset = MagicMock()
+    statefulset.spec.replicas = 2
+    statefulset.status.ready_replicas = 1
+
+    watcher = MagicMock()
+    watcher.stream = MagicMock()
+    watcher.stream.side_effect = urllib3.exceptions.ReadTimeoutError(
+        None, None, None)
+    watch.Watch.return_value = watcher
+
+    with pytest.raises(ActivityFailed) as excinfo:
+        statefulset_fully_available("mysvc")
+    assert "microservice 'mysvc' failed to recover within" in \
+        str(excinfo.value)
+
+
 @patch('chaosk8s.has_local_config_file', autospec=True)
 @patch('chaosk8s.pod.probes.client', autospec=True)
 @patch('chaosk8s.client')
