From 25590804cbe51c537f8ff0d493d0f785b92b94e8 Mon Sep 17 00:00:00 2001
From: Joshua K <joshua.michael.keyes@gmail.com>
Date: Fri, 30 Jul 2021 07:41:45 -0700
Subject: [PATCH] Add support for waiting on a StatefulSet. (#195)

Add support for waiting on a StatefulSet.

SUMMARY
This PR implements support for waiting on StatefulSet for readiness similar to how the other waiters currently work.
ISSUE TYPE

Feature Pull Request

ADDITIONAL INFORMATION
This was designed to (mostly) mimic the behaviour of the StatefulSetStatusViewer used by kubectl rollout status -w.

Reviewed-by: Abhijeet Kasurde <None>
Reviewed-by: Joshua K <None>
Reviewed-by: None <None>
Reviewed-by: Mike Graves <mgraves@redhat.com>
Reviewed-by: None <None>
---
 .../195-k8s-add-wait-statefulsets.yml         |   2 +
 molecule/default/tasks/waiter.yml             | 105 ++++++++++++++++++
 plugins/module_utils/common.py                |   9 ++
 3 files changed, 116 insertions(+)
 create mode 100644 changelogs/fragments/195-k8s-add-wait-statefulsets.yml

diff --git a/changelogs/fragments/195-k8s-add-wait-statefulsets.yml b/changelogs/fragments/195-k8s-add-wait-statefulsets.yml
new file mode 100644
index 00000000..8a334a87
--- /dev/null
+++ b/changelogs/fragments/195-k8s-add-wait-statefulsets.yml
@@ -0,0 +1,2 @@
+minor_changes:
+  - k8s - add support for waiting on statefulsets (https://github.com/ansible-collections/kubernetes.core/pull/195).
diff --git a/molecule/default/tasks/waiter.yml b/molecule/default/tasks/waiter.yml
index 44fc42b3..a2e5c1e5 100644
--- a/molecule/default/tasks/waiter.yml
+++ b/molecule/default/tasks/waiter.yml
@@ -140,6 +140,111 @@
           - ds.result.status.currentNumberScheduled == ds.result.status.desiredNumberScheduled
           - updated_ds_pods.resources[0].spec.containers[0].image.endswith(":3")
 
+    - name: Add a statefulset
+      k8s:
+        definition:
+          apiVersion: apps/v1
+          kind: StatefulSet
+          metadata:
+            name: wait-statefulset
+            namespace: "{{ wait_namespace }}"
+          spec:
+            selector:
+              matchLabels:
+                app: "{{ k8s_pod_name }}"
+            template: "{{ k8s_pod_template }}"
+        wait: yes
+        wait_sleep: 5
+        wait_timeout: 180
+      vars:
+        k8s_pod_name: wait-sts
+        k8s_pod_image: gcr.io/kuar-demo/kuard-amd64:1
+        k8s_pod_command:
+          - sleep
+          - "600"
+      register: sts
+
+    - name: Check that statefulset wait worked
+      assert:
+        that:
+          - sts.result.spec.replicas == sts.result.status.readyReplicas
+
+    - name: Update a statefulset in check_mode
+      k8s:
+        definition:
+          apiVersion: apps/v1
+          kind: StatefulSet
+          metadata:
+            name: wait-statefulset
+            namespace: "{{ wait_namespace }}"
+          spec:
+            selector:
+              matchLabels:
+                app: "{{ k8s_pod_name }}"
+            updateStrategy:
+              type: RollingUpdate
+            template: "{{ k8s_pod_template }}"
+        wait: yes
+        wait_sleep: 3
+        wait_timeout: 180
+      vars:
+        k8s_pod_name: wait-sts
+        k8s_pod_image: gcr.io/kuar-demo/kuard-amd64:2
+        k8s_pod_command:
+          - sleep
+          - "600"
+      register: update_sts_check_mode
+      check_mode: yes
+
+    - name: Check that check_mode result contains the changes
+      assert:
+        that:
+          - update_sts_check_mode is changed
+          - "update_sts_check_mode.result.spec.template.spec.containers[0].image == 'gcr.io/kuar-demo/kuard-amd64:2'"
+
+    - name: Update a statefulset
+      k8s:
+        definition:
+          apiVersion: apps/v1
+          kind: StatefulSet
+          metadata:
+            name: wait-statefulset
+            namespace: "{{ wait_namespace }}"
+          spec:
+            selector:
+              matchLabels:
+                app: "{{ k8s_pod_name }}"
+            updateStrategy:
+              type: RollingUpdate
+            template: "{{ k8s_pod_template }}"
+        wait: yes
+        wait_sleep: 3
+        wait_timeout: 180
+      vars:
+        k8s_pod_name: wait-sts
+        k8s_pod_image: gcr.io/kuar-demo/kuard-amd64:3
+        k8s_pod_command:
+          - sleep
+          - "600"
+      register: sts
+
+    - name: Get updated pods
+      k8s_info:
+        api_version: v1
+        kind: Pod
+        namespace: "{{ wait_namespace }}"
+        label_selectors:
+          - app=wait-sts
+        field_selectors:
+          - status.phase=Running
+      register: updated_sts_pods
+
+    - name: Check that statefulset wait worked
+      assert:
+        that:
+          - sts.result.spec.replicas == sts.result.status.readyReplicas
+          - updated_sts_pods.resources[0].spec.containers[0].image.endswith(":3")
+
     - name: Add a crashing pod
       k8s:
         definition:
diff --git a/plugins/module_utils/common.py b/plugins/module_utils/common.py
index cb090eff..de6989d3 100644
--- a/plugins/module_utils/common.py
+++ b/plugins/module_utils/common.py
@@ -401,6 +401,14 @@ class K8sAnsibleMixin(object):
                     and daemonset.status.observedGeneration == daemonset.metadata.generation
                     and not daemonset.status.unavailableReplicas)
 
+        def _statefulset_ready(statefulset):
+            return (statefulset.status and statefulset.spec.updateStrategy.type == "RollingUpdate"
+                    and statefulset.status.observedGeneration == (statefulset.metadata.generation or 0)
+                    and statefulset.status.updateRevision == statefulset.status.currentRevision
+                    and statefulset.status.updatedReplicas == statefulset.spec.replicas
+                    and statefulset.status.readyReplicas == statefulset.spec.replicas
+                    and statefulset.status.replicas == statefulset.spec.replicas)
+
         def _custom_condition(resource):
             if not resource.status or not resource.status.conditions:
                 return False
@@ -427,6 +435,7 @@ class K8sAnsibleMixin(object):
             return not resource or (resource.kind.endswith('List') and resource.items == [])
 
         waiter = dict(
+            StatefulSet=_statefulset_ready,
             Deployment=_deployment_ready,
             DaemonSet=_daemonset_ready,
             Pod=_pod_ready