From d47bd509191abc8217a5568a6c73324467355036 Mon Sep 17 00:00:00 2001
From: Tim Van Steenburgh <tvansteenburgh@gmail.com>
Date: Tue, 18 Oct 2016 09:37:26 -0400
Subject: [PATCH] Move add_unit logic into Application

---
 examples/relate.py  |  5 +++++
 juju/application.py | 22 ++++++++++++++++++++--
 juju/model.py       | 14 +++++---------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/examples/relate.py b/examples/relate.py
index 241d1a6..1952ffe 100644
--- a/examples/relate.py
+++ b/examples/relate.py
@@ -68,6 +68,11 @@ async def run():
         lambda delta, old_unit, new_unit, model:
             print('Unit removed: {}'.format(old_unit.entity_id))
     ))
+    unit_a, unit_b = await ubuntu_app.add_units(count=2)
+    unit_a.on_change(asyncio.coroutine(
+        lambda delta, old_unit, new_unit, model:
+            print('Unit changed: {}'.format(new_unit.entity_id))
+    ))
     await model.deploy(
         'nrpe',
         service_name='nrpe',
diff --git a/juju/application.py b/juju/application.py
index 53f584a..90a78ca 100644
--- a/juju/application.py
+++ b/juju/application.py
@@ -1,3 +1,4 @@
+import asyncio
 import logging
 
 from . import model
@@ -44,7 +45,7 @@ class Application(model.ModelEntity):
         """
         pass
 
-    def add_unit(self, count=1, to=None):
+    async def add_unit(self, count=1, to=None):
         """Add one or more units to this service.
 
         :param int count: Number of units to add
@@ -56,7 +57,24 @@ class Application(model.ModelEntity):
             If None, a new machine is provisioned.
 
         """
-        pass
+        app_facade = client.ApplicationFacade()
+        app_facade.connect(self.connection)
+
+        log.debug(
+            'Adding %s unit%s to %s',
+            count, '' if count == 1 else 's', self.name)
+
+        result = await app_facade.AddUnits(
+            application=self.name,
+            placement=to,
+            num_units=count,
+        )
+
+        return await asyncio.gather(*[
+            asyncio.ensure_future(self.model._wait_for_new('unit', unit_id))
+            for unit_id in result.units
+        ])
+
     add_units = add_unit
 
     def allocate(self, budget, value):
diff --git a/juju/model.py b/juju/model.py
index ab9f712..b168987 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -1295,16 +1295,12 @@ class BundleHandler(object):
             # doesn't, so we're not bothering, either
             unit_name = self._units_by_app[application].pop()
             log.debug('Reusing unit %s for %s', unit_name, application)
-            return unit_name
-        log.debug('Adding unit of %s%s',
-                  application,
-                  (' to %s' % placement) if placement else '')
-        result = await self.app_facade.AddUnits(
-            application=application,
-            placement=placement,
-            num_units=1,
+            return self.model.units[unit_name]
+
+        return await self.model.applications[application].add_unit(
+            count=1,
+            to=placement,
         )
-        return result.units[0]
 
     async def expose(self, application):
         """
-- 
2.25.1