Optimize & fix unit removal

juju/model.py

@@ -2157,32 +2157,35 @@ async def destroy_unit(self, unit_id, destroy_storage=False, dry_run=False, forc
         """Destroy units by name.
 
         """
-        connection = self.connection()
-        app_facade = client.ApplicationFacade.from_connection(connection)
-
-        # Get the corresponding unit tag
-        unit_tag = tag.unit(unit_id)
-        if unit_tag is None:
-            log.error("Error converting %s to a valid unit tag", unit_id)
-            return JujuUnitError("Error converting %s to a valid unit tag", unit_id)
-
-        log.debug(
-            'Destroying unit %s', unit_id)
-
-        return await app_facade.DestroyUnit(units=[{
-            'unit-tag': unit_tag,
-            'destroy-storage': destroy_storage,
-            'force': force,
-            'max-wait': max_wait,
-            'dry-run': dry_run,
-        }])
+        return await self.destroy_units(unit_id,
+                                        destroy_storage=destroy_storage,
+                                        dry_run=dry_run,
+                                        force=force,
+                                        max_wait=max_wait
+                                        )
 
     async def destroy_units(self, *unit_names, destroy_storage=False, dry_run=False, force=False, max_wait=None):
         """Destroy several units at once.
 
         """
-        for u in unit_names:
-            await self.destroy_unit(u, destroy_storage, dry_run, force, max_wait)
+        connection = self.connection()
+        app_facade = client.ApplicationFacade.from_connection(connection)
+
+        units_to_destroy = []
+        for unit_id in unit_names:
+            unit_tag = tag.unit(unit_id)
+            if unit_tag is None:
+                log.error("Error converting %s to a valid unit tag", unit_id)
+                raise JujuUnitError("Error converting %s to a valid unit tag", unit_id)
+            units_to_destroy.append({
+                'unit-tag': unit_tag,
+                'destroy-storage': destroy_storage,
+                'force': force,
+                'max-wait': max_wait,
+                'dry-run': dry_run,
+            })
+        log.debug('Destroying units %s', unit_names)
+        return await app_facade.DestroyUnit(units=units_to_destroy)
 
     def download_backup(self, archive_id, target_filename=None):
         """Download a backup archive file.

juju/unit.py

@@ -102,7 +102,7 @@ def get_subordinates(self):
         return [u for u_name, u in self.model.units.items() if u.is_subordinate and
                 u.principal_unit == self.name]
 
-    async def destroy(self):
+    async def destroy(self, destroy_storage=False, dry_run=False, force=False, max_wait=None):
         """Destroy this unit.
 
         """
@@ -111,7 +111,12 @@ async def destroy(self):
         log.debug(
             'Destroying %s', self.name)
 
-        return await app_facade.DestroyUnit(units=[{"unit-tag": self.name}])
+        return await app_facade.DestroyUnit(units=[{"unit-tag": self.tag,
+                                                    'destroy-storage': destroy_storage,
+                                                    'force': force,
+                                                    'max-wait': max_wait,
+                                                    'dry-run': dry_run,
+                                                    }])
     remove = destroy
 
     async def get_public_address(self):

tests/integration/test_unit.py

@@ -280,3 +280,19 @@ async def test_subordinate_units(event_loop):
         assert n_unit.principal_unit == 'ubuntu/0'
         assert u_unit.principal_unit == ''
         assert [u.name for u in u_unit.get_subordinates()] == [n_unit.name]
+
+
+@base.bootstrapped
+async def test_destroy_unit(event_loop):
+    async with base.CleanModel() as model:
+        app = await model.deploy(
+            'juju-qa-test',
+            application_name='test',
+            num_units=3,
+        )
+        # wait for the units to come up
+        await model.block_until(lambda: app.units, timeout=480)
+
+        await app.units[0].destroy()
+        await asyncio.sleep(5)
+        assert len(app.units) == 2