From c8ec5f419a5146ea92ab440b4839b392b2bb6d65 Mon Sep 17 00:00:00 2001
From: Daniel Turner <daniel.turner@shopify.com>
Date: Wed, 16 Dec 2020 15:07:25 -0800
Subject: [PATCH] Track in and out

---
 lib/krane/deploy_task.rb                  | 14 +++++++++-----
 lib/krane/resource_deployer.rb            |  4 +++-
 test/unit/krane/resource_deployer_test.rb |  2 +-
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/lib/krane/deploy_task.rb b/lib/krane/deploy_task.rb
index 9ef8473d9..ba9557633 100644
--- a/lib/krane/deploy_task.rb
+++ b/lib/krane/deploy_task.rb
@@ -57,7 +57,7 @@ class DeployTask
     )
 
     def predeploy_sequence
-      default_group = { skip_groups: [] }
+      default_group = { groups: [], skip_groups: [] }
       before_crs = %w(
         ResourceQuota
         NetworkPolicy
@@ -71,10 +71,14 @@ def predeploy_sequence
 
       after_crs = %w(
         Pod
-      ).map { |r| [r, default_group] }
+      )
+
+      predeploy_hash = Hash[before_crs]
+      cluster_resource_discoverer.crds.select(&:predeployed?).each do |cr|
+        predeploy_hash[cr.kind] ||= { groups: [cr.group], skip_groups: [] }
+      end
 
-      crs = cluster_resource_discoverer.crds.select(&:predeployed?).map { |cr| [cr.kind, default_group] }
-      predeploy_hash = Hash[before_crs + crs + after_crs]
+      after_crs.each { |cr| predeploy_hash[cr] = default_group }
 
       cluster_resource_discoverer.crds.reject(&:predeployed?).each do |cr|
         predeploy_hash[cr.kind][:skip_groups] << cr.group if predeploy_hash[cr.kind]
@@ -222,7 +226,7 @@ def ejson_provisioners
     def deploy_has_priority_resources?(resources)
       resources.any? do |r|
         next unless (pr = predeploy_sequence[r.type])
-        !pr[:skip_groups].include?(r.group)
+        (pr[:groups].empty? || pr[:groups].include?(r.group)) && !pr[:skip_groups].include?(r.group)
       end
     end
 
diff --git a/lib/krane/resource_deployer.rb b/lib/krane/resource_deployer.rb
index b0cbc7782..0d9eed4c1 100644
--- a/lib/krane/resource_deployer.rb
+++ b/lib/krane/resource_deployer.rb
@@ -48,7 +48,9 @@ def predeploy_priority_resources(resource_list, predeploy_sequence)
 
       predeploy_sequence.each do |resource_type, attributes|
         matching_resources = resource_list.select do |r|
-          r.type == resource_type && !attributes[:skip_groups].include?(r.group)
+          r.type == resource_type &&
+          ((attributes[:groups].empty? || attributes[:groups].include?(r.group)) &&
+            !attributes[:skip_groups].include?(r.group))
         end
         StatsD.client.gauge('priority_resources.count', matching_resources.size, tags: statsd_tags)
 
diff --git a/test/unit/krane/resource_deployer_test.rb b/test/unit/krane/resource_deployer_test.rb
index 7c15e56c7..31a040918 100644
--- a/test/unit/krane/resource_deployer_test.rb
+++ b/test/unit/krane/resource_deployer_test.rb
@@ -71,7 +71,7 @@ def test_predeploy_priority_resources_respects_pre_deploy_list
     # are deployed. See test_predeploy_priority_resources_respects_empty_pre_deploy_list
     # for counter example
     Krane::ResourceWatcher.expects(:new).returns(watcher)
-    priority_list = { kind => { group: "core" } }
+    priority_list = { kind => { groups: ["core"], skip_groups: [] } }
     resource_deployer.predeploy_priority_resources([resource], priority_list)
   end