From fe5e5629a326f8a87aa3dc12283a914d8597cf02 Mon Sep 17 00:00:00 2001
From: "Joshua M. Keyes" <joshua.michael.keyes@gmail.com>
Date: Thu, 29 Jul 2021 12:01:03 -0700
Subject: [PATCH] Add support for waiting on a StatefulSet.

---
 .../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 0000000000..8a334a873a
--- /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 44fc42b3ff..a2e5c1e5aa 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 70c6fb19e0..e4b4aac220 100644
--- a/plugins/module_utils/common.py
+++ b/plugins/module_utils/common.py
@@ -397,6 +397,14 @@ def _daemonset_ready(daemonset):
                     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
@@ -423,6 +431,7 @@ def _resource_absent(resource):
             return not resource
 
         waiter = dict(
+            StatefulSet=_statefulset_ready,
             Deployment=_deployment_ready,
             DaemonSet=_daemonset_ready,
             Pod=_pod_ready