From f93420deb029b6938f73d28a4f7a414b2b9655dc Mon Sep 17 00:00:00 2001
From: Pete Vander Giessen <petevg@gmail.com>
Date: Wed, 5 Apr 2017 17:46:24 -0400
Subject: [PATCH] Added Monitor class to Connection. (#105)

Added Monitor class to Connection.

Calling Monitor.status allows a client to determine whether the
connection is "connected", "disconnected" or in an "error" state, and
take appropriate action as a result.

This is meant to address the issue of watchers failing silently when the
websocket connection goes away due to a network issue.
---
 juju/client/connection.py            | 82 +++++++++++++++++++++++++++-
 tests/integration/test_connection.py | 28 ++++++++++
 2 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/juju/client/connection.py b/juju/client/connection.py
index 3ee8f16..6c31ab6 100644
--- a/juju/client/connection.py
+++ b/juju/client/connection.py
@@ -21,6 +21,83 @@ from juju.utils import IdQueue
 log = logging.getLogger("websocket")
 
 
+class Monitor:
+    """
+    Monitor helper class for our Connection class.
+
+    Contains a reference to an instantiated Connection, along with a
+    reference to the Connection.receiver Future. Upon inspecttion of
+    these objects, this class determines whether the connection is in
+    an 'error', 'connected' or 'disconnected' state.
+
+    Use this class to stay up to date on the health of a connection,
+    and take appropriate action if the connection errors out due to
+    network issues or other unexpected circumstances.
+
+    """
+    ERROR = 'error'
+    CONNECTED = 'connected'
+    DISCONNECTED = 'disconnected'
+    UNKNOWN = 'unknown'
+
+    def __init__(self, connection):
+        self.connection = connection
+        self.receiver = None
+
+    @property
+    def status(self):
+        """
+        Determine the status of the connection and receiver, and return
+        ERROR, CONNECTED, or DISCONNECTED as appropriate.
+
+        For simplicity, we only consider ourselves to be connected
+        after the Connection class has setup a receiver task. This
+        only happens after the websocket is open, and the connection
+        isn't usable until that receiver has been started.
+
+        """
+
+        # DISCONNECTED: connection not yet open
+        if not self.connection.ws:
+            return self.DISCONNECTED
+        if not self.receiver:
+            return self.DISCONNECTED
+
+        # ERROR: Connection closed (or errored), but we didn't call
+        # connection.close
+        if not self.connection.close_called and self.receiver_exceptions():
+            return self.ERROR
+        if not self.connection.close_called and not self.connection.ws.open:
+            # The check for self.receiver existing above guards against the
+            # case where we're not open because we simply haven't
+            # setup the connection yet.
+            return self.ERROR
+
+        # DISCONNECTED: cleanly disconnected.
+        if self.connection.close_called and not self.connection.ws.open:
+            return self.DISCONNECTED
+
+        # CONNECTED: everything is fine!
+        if self.connection.ws.open:
+            return self.CONNECTED
+
+        # UNKNOWN: We should never hit this state -- if we do,
+        # something went wrong with the logic above, and we do not
+        # know what state the connection is in.
+        return self.UNKNOWN
+
+    def receiver_exceptions(self):
+        """
+        Return exceptions in the receiver, if any.
+
+        """
+        if not self.receiver:
+            return None
+        if not self.receiver.done():
+            return None
+        return self.receiver.exception()
+
+
 class Connection:
     """
     Usage::
@@ -54,6 +131,8 @@ class Connection:
         self.ws = None
         self.facades = {}
         self.messages = IdQueue(loop=self.loop)
+        self.close_called = False
+        self.monitor = Monitor(connection=self)
 
     @property
     def is_open(self):
@@ -76,11 +155,12 @@ class Connection:
         kw['loop'] = self.loop
         self.addr = url
         self.ws = await websockets.connect(url, **kw)
-        self.loop.create_task(self.receiver())
+        self.monitor.receiver = self.loop.create_task(self.receiver())
         log.info("Driver connected to juju %s", url)
         return self
 
     async def close(self):
+        self.close_called = True
         await self.ws.close()
 
     async def recv(self, request_id):
diff --git a/tests/integration/test_connection.py b/tests/integration/test_connection.py
index 9c61759..18c76b4 100644
--- a/tests/integration/test_connection.py
+++ b/tests/integration/test_connection.py
@@ -12,3 +12,31 @@ async def test_connect_current(event_loop):
 
         assert isinstance(conn, Connection)
         await conn.close()
+
+
+@base.bootstrapped
+@pytest.mark.asyncio
+async def test_monitor(event_loop):
+
+    async with base.CleanModel():
+        conn = await Connection.connect_current()
+
+        assert conn.monitor.status == 'connected'
+        await conn.close()
+
+        assert conn.monitor.status == 'disconnected'
+
+
+@base.bootstrapped
+@pytest.mark.asyncio
+async def test_monitor_catches_error(event_loop):
+
+    async with base.CleanModel():
+        conn = await Connection.connect_current()
+
+        assert conn.monitor.status == 'connected'
+        await conn.ws.close()
+
+        assert conn.monitor.status == 'error'
+
+        await conn.close()
-- 
2.17.1