From d571b9bbbc09c937e83b5ca567a4507b3eb806b6 Mon Sep 17 00:00:00 2001
From: Alan Rominger <arominge@redhat.com>
Date: Wed, 2 Nov 2022 15:14:12 -0400
Subject: [PATCH] Refactor test_get_cleanup_task_kwargs_active_jobs and add new
 test

This takes some logic out of the queryset logic,
  using some established assumptions about the task manager
  if a job lands on a hybrid node (or is a project update) then
  it will have the same controller and execution node

With that established, the queryset can be simplified
---
 awx/main/models/ha.py                       | 11 ++++++-----
 awx/main/tests/functional/test_instances.py | 20 +++++++++++++++++++-
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/awx/main/models/ha.py b/awx/main/models/ha.py
index dbeb81dcac..f101a94d7a 100644
--- a/awx/main/models/ha.py
+++ b/awx/main/models/ha.py
@@ -233,11 +233,12 @@ class Instance(HasPolicyEditsMixin, BaseModel):
         if not isinstance(vargs.get('grace_period'), int):
             vargs['grace_period'] = 60  # grace period of 60 minutes, need to set because CLI default will not take effect
         if 'exclude_strings' not in vargs and vargs.get('file_pattern'):
-            active_pks = list(
-                UnifiedJob.objects.filter(
-                    (models.Q(execution_node=self.hostname) | models.Q(controller_node=self.hostname)) & models.Q(status__in=('running', 'waiting'))
-                ).values_list('pk', flat=True)
-            )
+            active_job_qs = UnifiedJob.objects.filter(status__in=('running', 'waiting'))
+            if self.node_type == 'execution':
+                active_job_qs = active_job_qs.filter(execution_node=self.hostname)
+            else:
+                active_job_qs = active_job_qs.filter(controller_node=self.hostname)
+            active_pks = list(active_job_qs.values_list('pk', flat=True))
             if active_pks:
                 vargs['exclude_strings'] = [JOB_FOLDER_PREFIX % job_id for job_id in active_pks]
         if 'remove_images' in vargs or 'image_prune' in vargs:
diff --git a/awx/main/tests/functional/test_instances.py b/awx/main/tests/functional/test_instances.py
index 8ce6524d38..df6d177868 100644
--- a/awx/main/tests/functional/test_instances.py
+++ b/awx/main/tests/functional/test_instances.py
@@ -1,7 +1,7 @@
 import pytest
 from unittest import mock
 
-from awx.main.models import AdHocCommand, InventoryUpdate, JobTemplate
+from awx.main.models import AdHocCommand, InventoryUpdate, JobTemplate, Job
 from awx.main.models.activity_stream import ActivityStream
 from awx.main.models.ha import Instance, InstanceGroup
 from awx.main.tasks.system import apply_cluster_membership_policies
@@ -15,6 +15,24 @@ def test_default_tower_instance_group(default_instance_group, job_factory):
     assert default_instance_group in job_factory().preferred_instance_groups
 
 
+@pytest.mark.django_db
+@pytest.mark.parametrize('node_type', ('execution', 'control'))
+@pytest.mark.parametrize('active', (True, False))
+def test_get_cleanup_task_kwargs_active_jobs(node_type, active):
+    instance = Instance.objects.create(hostname='foobar', node_type=node_type)
+    job_kwargs = dict()
+    job_kwargs['controller_node' if node_type == 'control' else 'execution_node'] = instance.hostname
+    job_kwargs['status'] = 'running' if active else 'successful'
+
+    job = Job.objects.create(**job_kwargs)
+    kwargs = instance.get_cleanup_task_kwargs()
+
+    if active:
+        assert kwargs['exclude_strings'] == [f'awx_{job.pk}_']
+    else:
+        assert 'exclude_strings' not in kwargs
+
+
 @pytest.mark.django_db
 class TestPolicyTaskScheduling:
     """Tests make assertions about when the policy task gets scheduled"""