Merge "Exclude generated objects from wait logic"

armada/handlers/wait.py

@@ -154,15 +154,23 @@ class ResourceWait(ABC):
         '''
         pass
 
-    def include_resource(self, resource):
+    def get_exclude_reason(self, resource):
         '''
-        Test to include or exclude a resource in a wait operation. This method
+        If a resource should be excluded from the wait, returns a message
-        can be used to exclude resources that should not be included in wait
+        to explain why. Unless overridden, all resources are included.
-        operations (e.g. test pods).
         :param resource: resource to test
-        :returns: boolean representing test result
+        :returns: string representing exclude reason
         '''
-        return True
+        return None
+
+    def include_resource(self, resource):
+        exclude_reason = self.get_exclude_reason(resource)
+
+        if exclude_reason:
+            LOG.debug('Excluding %s %s from wait: %s', self.resource_type,
+                      resource.metadata.name, exclude_reason)
+
+        return not exclude_reason
 
     def handle_resource(self, resource):
         resource_name = resource.metadata.name
@@ -353,18 +361,21 @@ class PodWait(ResourceWait):
             resource_type, chart_wait, labels,
             chart_wait.k8s.client.list_namespaced_pod, **kwargs)
 
-    def include_resource(self, resource):
+    def get_exclude_reason(self, resource):
         pod = resource
-        include = not is_test_pod(pod)
 
-        # NOTE(drewwalters96): Test pods may cause wait operations to fail
+        # Exclude helm test pods
-        # when old resources remain from previous upgrades/tests. Indicate that
+        # TODO: Possibly exclude other helm hook pods/jobs (besides tests)?
-        # test pods should not be included in wait operations.
+        if is_test_pod(pod):
-        if not include:
+            return 'helm test pod'
-            LOG.debug('Test pod %s will be skipped during wait operations.',
-                      pod.metadata.name)
 
-        return include
+        # Exclude job pods
+        # TODO: Once controller-based waits are enabled by default, ignore
+        # controller-owned pods as well.
+        if has_owner(pod, 'Job'):
+            return 'generated by job (wait on that instead if not already)'
+
+        return None
 
     def is_resource_ready(self, resource):
         pod = resource
@@ -392,6 +403,15 @@ class JobWait(ResourceWait):
             resource_type, chart_wait, labels,
             chart_wait.k8s.batch_api.list_namespaced_job, **kwargs)
 
+    def get_exclude_reason(self, resource):
+        job = resource
+
+        # Exclude cronjob jobs
+        if has_owner(job, 'CronJob'):
+            return 'generated by cronjob (not part of release)'
+
+        return None
+
     def is_resource_ready(self, resource):
         job = resource
         name = job.metadata.name
@@ -406,6 +426,16 @@ class JobWait(ResourceWait):
         return (msg.format(name), True)
 
 
+def has_owner(resource, kind=None):
+    owner_references = resource.metadata.owner_references or []
+
+    for owner in owner_references:
+        if not kind or kind == owner.kind:
+            return True
+
+    return False
+
+
 CountOrPercent = collections.namedtuple('CountOrPercent',
                                         'number is_percent source')
 

armada/tests/unit/handlers/test_wait.py

@@ -194,12 +194,13 @@ class PodWaitTestCase(base.ArmadaTestCase):
 
     def test_include_resource(self):
 
-        def mock_resource(annotations):
+        def mock_resource(annotations={}, owner_references=None):
             resource = mock.Mock()
             resource.metadata.annotations = annotations
+            resource.metadata.owner_references = owner_references
             return resource
 
-        test_resources = [
+        test_pods = [
             mock_resource({
                 'key': 'value',
                 'helm.sh/hook': 'test-success'
@@ -207,18 +208,69 @@ class PodWaitTestCase(base.ArmadaTestCase):
             mock_resource({'helm.sh/hook': 'test-failure'}),
             mock_resource({'helm.sh/hook': 'test-success,pre-install'})
         ]
-        non_test_resources = [
+        job_pods = [
+            mock_resource(owner_references=[mock.Mock(kind='Job')]),
+            mock_resource(owner_references=[
+                mock.Mock(kind='NotAJob'),
+                mock.Mock(kind='Job')
+            ])
+        ]
+        included_pods = [
+            mock_resource(),
+            mock_resource(owner_references=[]),
             mock_resource({'helm.sh/hook': 'pre-install'}),
             mock_resource({'key': 'value'}),
-            mock_resource({})
+            mock_resource(owner_references=[mock.Mock(kind='NotAJob')])
         ]
 
         unit = self.get_unit({})
 
-        # Validate test resources excluded
+        # Validate test pods excluded
-        for resource in test_resources:
+        for pod in test_pods:
-            self.assertFalse(unit.include_resource(resource))
+            self.assertFalse(unit.include_resource(pod))
+
+        # Validate test pods excluded
+        for pod in job_pods:
+            self.assertFalse(unit.include_resource(pod))
 
         # Validate other resources included
-        for resource in non_test_resources:
+        for pod in included_pods:
-            self.assertTrue(unit.include_resource(resource))
+            self.assertTrue(unit.include_resource(pod))
+
+
+class JobWaitTestCase(base.ArmadaTestCase):
+
+    def get_unit(self, labels):
+        return wait.JobWait(
+            resource_type='job', chart_wait=mock.MagicMock(), labels=labels)
+
+    def test_include_resource(self):
+
+        def mock_resource(annotations={}, owner_references=None):
+            resource = mock.Mock()
+            resource.metadata.annotations = annotations
+            resource.metadata.owner_references = owner_references
+            return resource
+
+        cronjob_jobs = [
+            mock_resource(owner_references=[mock.Mock(kind='CronJob')]),
+            mock_resource(owner_references=[
+                mock.Mock(kind='NotACronJob'),
+                mock.Mock(kind='CronJob')
+            ])
+        ]
+        included_jobs = [
+            mock_resource(),
+            mock_resource(owner_references=[]),
+            mock_resource(owner_references=[mock.Mock(kind='NotAJob')])
+        ]
+
+        unit = self.get_unit({})
+
+        # Validate test pods excluded
+        for job in cronjob_jobs:
+            self.assertFalse(unit.include_resource(job))
+
+        # Validate other resources included
+        for job in included_jobs:
+            self.assertTrue(unit.include_resource(job))