correct capacity algorithm for task manager

awx/main/tests/functional/task_management/test_scheduler.py

@@ -4,15 +4,43 @@ from datetime import timedelta, datetime
 
 from django.core.cache import cache
 from django.utils.timezone import now as tz_now
+from django.test import TransactionTestCase
 
 from awx.main.scheduler import TaskManager
 from awx.main.models import (
     Job,
     Instance,
+    InstanceGroup,
     WorkflowJob,
 )
 
 
+
+@pytest.mark.django_db
+class TestCapacityMapping(TransactionTestCase):
+
+    def sample_cluster(self):
+        ig_small = InstanceGroup.objects.create(name='ig_small')
+        ig_large = InstanceGroup.objects.create(name='ig_large')
+        tower = InstanceGroup.objects.create(name='tower')
+        i1 = Instance.objects.create(hostname='i1', capacity=200)
+        i2 = Instance.objects.create(hostname='i2', capacity=200)
+        i3 = Instance.objects.create(hostname='i3', capacity=200)
+        ig_small.instances.add(i1)
+        ig_large.instances.add(i2, i3)
+        tower.instances.add(i2)
+
+    def test_something(self):
+        self.sample_cluster()
+        with self.assertNumQueries(2):
+            inst_map, ig_map = Instance.objects.capacity_mapping()
+        assert inst_map['i1'] == set(['ig_small'])
+        assert inst_map['i2'] == set(['ig_large', 'tower'])
+        assert ig_map['ig_small'] == set(['ig_small'])
+        assert ig_map['ig_large'] == set(['ig_large', 'tower'])
+        assert ig_map['tower'] == set(['ig_large', 'tower'])
+
+
 @pytest.mark.django_db
 def test_single_job_scheduler_launch(default_instance_group, job_template_factory, mocker):
     objects = job_template_factory('jt', organization='org1', project='proj',