Fixed issues introduced during testing

daskos · Jan 30, 2017 · 607f4d7 · 607f4d7
1 parent e2b579f
commit 607f4d7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/mentos/subscription.py b/mentos/subscription.py
@@ -226,7 +226,8 @@ def _event_handler(self, message):
                 if "framework_id" in message["subscribed"]:
                     self.framework["id"] = message[
                         "subscribed"]["framework_id"]
-                if self.state.current_state==States.SUBSCRIBING:
+                # Add special check to ensure subscribed for executor
+                if self.state.current_state==States.SUBSCRIBING and "executor_info" in message["subscribed"]:
                     self.state.transition_to(States.SUBSCRIBED)
 
             if message["type"] in self.event_handlers:

diff --git a/mentos/tests/test_subscription.py b/mentos/tests/test_subscription.py
@@ -29,10 +29,10 @@ def test_subscription(io_loop, mocker):
     assert sub.state.current_state == states.States.CLOSED
     yield sub.start()
 
-    yield sub.ensure_safe([states.States.SUBSCRIBING])
+    yield sub.ensure_safe([states.States.SUBSCRIBING,states.States.SUBSCRIBED])
 
 
-    assert sub.state.current_state == states.States.SUBSCRIBING
+    assert sub.state.current_state in [states.States.SUBSCRIBING,states.States.SUBSCRIBED]
 
     assert "framework_id" not in sub.framework
     assert sub.connection != None