update tests

awx/main/tests/functional/tasks/test_tasks_system.py

@@ -8,9 +8,16 @@ from unittest import mock
 
 import pytest
 
-from awx.main.tasks.system import CleanupImagesAndFiles, execution_node_health_check, inspect_established_receptor_connections, clear_setting_cache
+from awx.main.tasks.system import (
+    CleanupImagesAndFiles,
+    execution_node_health_check,
+    inspect_established_receptor_connections,
+    clear_setting_cache,
+    _batched_delete_inventory,
+)
 from awx.main.management.commands.dispatcherd import Command
-from awx.main.models import Instance, Job, ReceptorAddress, InstanceLink
+from awx.main.models import Instance, Inventory, Job, Organization, ReceptorAddress, InstanceLink
+from awx.main.models.inventory import Group, Host
 
 
 @pytest.mark.django_db
@@ -105,6 +112,39 @@ def test_folder_cleanup_multiple_running_jobs(job_folder_factory, me_inst):
     assert [os.path.exists(d) for d in dirs] == [True for i in range(num_jobs)]
 
 
+@pytest.mark.django_db
+class TestBatchedDeleteInventory:
+    def _make_inventory_with_hosts(self, count):
+        org = Organization.objects.create(name='test-org')
+        inv = Inventory.objects.create(name='test-inv', organization=org)
+        group = Group.objects.create(name='test-group', inventory=inv)
+        hosts = [Host(name=f'host-{i}', inventory=inv) for i in range(count)]
+        Host.objects.bulk_create(hosts)
+        group.hosts.set(Host.objects.filter(inventory=inv))
+        return inv
+
+    def test_deletes_all_hosts_and_inventory(self):
+        inv = self._make_inventory_with_hosts(10)
+        inv_id = inv.id
+        _batched_delete_inventory(inv, batch_size=3)
+        assert not Host.objects.filter(inventory_id=inv_id).exists()
+        assert not Group.objects.filter(inventory_id=inv_id).exists()
+        assert not Inventory.objects.filter(id=inv_id).exists()
+
+    def test_no_hosts(self):
+        inv = self._make_inventory_with_hosts(0)
+        inv_id = inv.id
+        _batched_delete_inventory(inv)
+        assert not Inventory.objects.filter(id=inv_id).exists()
+
+    def test_exactly_one_batch(self):
+        inv = self._make_inventory_with_hosts(5)
+        inv_id = inv.id
+        _batched_delete_inventory(inv, batch_size=5)
+        assert not Host.objects.filter(inventory_id=inv_id).exists()
+        assert not Inventory.objects.filter(id=inv_id).exists()
+
+
 @pytest.mark.django_db
 def test_clear_setting_cache_log_level_branch(settings):
     settings.LOG_AGGREGATOR_LEVEL = 'DEBUG'