32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
from django.test import TransactionTestCase
|
|
|
|
from awx.main.models import (
|
|
Instance,
|
|
InstanceGroup,
|
|
)
|
|
|
|
|
|
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)
|
|
return [tower, ig_large, ig_small]
|
|
|
|
def test_mapping(self):
|
|
self.sample_cluster()
|
|
with self.assertNumQueries(2):
|
|
inst_map, ig_map = InstanceGroup.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'])
|