X-Git-Url: https://osm.etsi.org/gitweb/?p=osm%2FN2VC.git;a=blobdiff_plain;f=juju%2Fmodel.py;h=1ab29b5ca8f920c9d6b24183b9baab8b0da78760;hp=ef6bc85115c3e213938bebec3c8fce31cf07263d;hb=e2a3f01e7bb983ac3cd8eb24117e0d4208b40cc9;hpb=21ba8065355c663597df7bd2b481ade519f89a79

diff --git a/juju/model.py b/juju/model.py
index ef6bc85..1ab29b5 100644
--- a/juju/model.py
+++ b/juju/model.py
@@ -671,7 +671,7 @@ class Model(object):
 
         :param entity_type: The entity's type.
         :param entity_id: The entity's id.
-        :param action: the type of action (e.g., 'add' or 'change')
+        :param action: the type of action (e.g., 'add', 'change', or 'remove')
         :param predicate: optional callable that must take as an
             argument a delta, and must return a boolean, indicating
             whether the delta contains the specific action we're looking
@@ -686,7 +686,9 @@ class Model(object):
 
         self.add_observer(callback, entity_type, action, entity_id, predicate)
         entity_id = await q.get()
-        return self.state._live_entity_map(entity_type)[entity_id]
+        # object might not be in the entity_map if we were waiting for a
+        # 'remove' action
+        return self.state._live_entity_map(entity_type).get(entity_id)
 
     async def _wait_for_new(self, entity_type, entity_id=None, predicate=None):
         """Wait for a new object to appear in the Model and return it.