Potential checks for OnDelete strategy with full and none annotations

Shopify · Jun 6, 2023 · 0ba7a26 · 0ba7a26
1 parent 3e97f0f
commit 0ba7a26
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 14 deletions.
diff --git a/lib/krane/kubernetes_resource/stateful_set.rb b/lib/krane/kubernetes_resource/stateful_set.rb
@@ -5,6 +5,9 @@ class StatefulSet < PodSetBase
     TIMEOUT = 10.minutes
     ONDELETE = 'OnDelete'
     SYNC_DEPENDENCIES = %w(Pod)
+    REQUIRED_ROLLOUT_ANNOTATION = "required-rollout"
+    REQUIRED_ROLLOUT_TYPES = %w(full none).freeze
+    DEFAULT_REQUIRED_ROLLOUT = 'full'
     attr_reader :pods
 
     def sync(cache)
@@ -19,20 +22,27 @@ def status
     end
 
     def deploy_succeeded?
-      success = observed_generation == current_generation &&
-        desired_replicas == status_data['readyReplicas'].to_i &&
-        desired_replicas == status_data['currentReplicas'].to_i
+      success = observed_generation == current_generation
+
       if update_strategy == ONDELETE
-        success &= desired_replicas == status_data['updatedReplicas'].to_i
-        # Gem cannot monitor update since it doesn't occur until delete
         unless @success_assumption_warning_shown
           @logger.warn("WARNING: Your StatefulSet's updateStrategy is set to OnDelete, "\
                        "which means updates will not be applied until its pods are deleted.")
           @success_assumption_warning_shown = true
         end
+
+        if required_rollout == 'none'
+          return true
+        elsif required_rollout == 'full'
+          success &= desired_replicas == status_data['updatedReplicas'].to_i
+        end
       else
         success &= status_data['currentRevision'] == status_data['updateRevision']
       end
+
+      success &= desired_replicas == status_data['readyReplicas'].to_i &&
+        desired_replicas == status_data['currentReplicas'].to_i
+
       success
     end
 
@@ -66,5 +76,9 @@ def parent_of_pod?(pod_data)
       pod_data["metadata"]["ownerReferences"].any? { |ref| ref["uid"] == @instance_data["metadata"]["uid"] } &&
       @instance_data["status"]["currentRevision"] == pod_data["metadata"]["labels"]["controller-revision-hash"]
     end
+
+    def required_rollout
+      krane_annotation_value("required-rollout") || DEFAULT_REQUIRED_ROLLOUT
+    end
   end
 end
diff --git a/test/fixtures/for_unit_tests/stateful_set_test.yml b/test/fixtures/for_unit_tests/stateful_set_test.yml
@@ -3,6 +3,8 @@ kind: StatefulSet
 metadata:
   name: test-ss
   generation: 2
+  annotations:
+    krane.shopify.io/required-rollout: full
   labels:
     app: hello-cloud
     name: test-ss

diff --git a/test/unit/krane/kubernetes_resource/stateful_set_test.rb b/test/unit/krane/kubernetes_resource/stateful_set_test.rb
@@ -4,26 +4,28 @@
 class StatefulSetTest < Krane::TestCase
   include ResourceCacheTestHelper
 
-  def test_deploy_succeeded_is_true_when_revision_and_replica_counts_match_for_rollingupdate_strategy
+  def test_deploy_succeeded_when_revision_and_replica_counts_match_for_rollingupdate_strategy
     template = build_ss_template(status: { "observedGeneration": 2 }, updateStrategy: "RollingUpdate")
     ss = build_synced_ss(template: template)
     assert_predicate(ss, :deploy_succeeded?)
   end
 
-  def test_deploy_succeeded_is_true_when_revision_and_replica_counts_match_for_ondelete_strategy
+  def test_deploy_succeeded_when_replica_counts_match_for_ondelete_strategy_with_full_annotation
     template = build_ss_template(status: { "observedGeneration": 2 }, updateStrategy: "OnDelete")
     ss = build_synced_ss(template: template)
     assert_predicate(ss, :deploy_succeeded?)
   end
 
-  def test_deploy_failed_ensures_controller_has_observed_deploy_for_rollingupdate_strategy
-    template = build_ss_template(status: { "observedGeneration": 1 }, updateStrategy: "RollingUpdate")
+  def test_deploy_succeeded_when_latest_rs_created_for_ondelete_strategy_with_none_annotation
+    template = build_ss_template(status: { "observedGeneration": 2 }, updateStrategy: "OnDelete", rollout: "none")
     ss = build_synced_ss(template: template)
-    refute_predicate(ss, :deploy_succeeded?)
+    assert_predicate(ss, :deploy_succeeded?)
   end
 
-  def test_deploy_failed_ensures_controller_has_observed_deploy_for_ondelete_strategy
-    template = build_ss_template(status: { "observedGeneration": 1 }, updateStrategy: "OnDelete")
+  # add test - that deploy fails if replica counts aren't equal
+
+  def test_deploy_failed_ensures_controller_has_observed_deploy_for_rollingupdate_strategy
+    template = build_ss_template(status: { "observedGeneration": 1 }, updateStrategy: "RollingUpdate")
     ss = build_synced_ss(template: template)
     refute_predicate(ss, :deploy_succeeded?)
   end
@@ -52,8 +54,12 @@ def test_deploy_failed_not_fooled_by_stale_status_for_ondelete_strategy
 
   private
 
-  def build_ss_template(status: {}, updateStrategy:)
-    ss_fixture.dup.deep_merge("status" => status, "spec" => {"updateStrategy" => {"type" => updateStrategy}})
+  def build_ss_template(status: {}, updateStrategy: "RollingUpdate", rollout: "full")
+    ss_fixture.dup.deep_merge(
+      "status" => status,
+      "spec" => {"updateStrategy" => {"type" => updateStrategy}},
+      "metadata" => {"annotations" => {"krane.shopify.io/rollout" => rollout}}
+    )
   end
 
   def build_synced_ss(template:)