Fixed auto-reconnect for proxy and created tests for auto-reconnect

tests/environment/__init__.py

@@ -65,7 +65,12 @@ class TestEnvironment(object):
         # Destroy the env in case it wasn't killed correctly last time.
         subprocess.check_call(["docker", "compose", "down"], stdout=subprocess.DEVNULL, cwd=thisdir)
         self._subp = _docker_process = subprocess.Popen(["docker", "compose", "up", "--build", "--force-recreate", "--no-deps"], stdout=subprocess.DEVNULL, cwd=thisdir)
-        timeout = 30
+        if not self._wait_for_meshcentral():
+            self.__exit__(None, None, None)
+            raise Exception("Failed to create docker instance")
+        return self
+
+    def _wait_for_meshcentral(self, timeout=30):
         start = time.time()
         while time.time() - start < timeout:
             try:
@@ -82,10 +87,8 @@ class TestEnvironment(object):
                 pass
             time.sleep(1)
         else:
-            self.__exit__(None, None, None)
-            raise Exception("Failed to create docker instance")
-        return self
-
+            return False
+        return True
 
     def __exit__(self, exc_t, exc_v, exc_tb):
         pass
@@ -93,6 +96,15 @@ class TestEnvironment(object):
     def create_agent(self, meshid):
         return Agent(meshid, self.mcurl, self.clienturl, self.dockerurl)
 
+    # Restart our docker instances, to test reconnect code.
+    def restart_mesh(self):
+        subprocess.check_call(["docker", "container", "restart", "meshctrl-meshcentral"], stdout=subprocess.DEVNULL, cwd=thisdir)
+        assert self._wait_for_meshcentral(), "Failed to restart docker instance"
+
+    def restart_proxy(self):
+        subprocess.check_call(["docker", "container", "restart", "meshctrl-squid"], stdout=subprocess.DEVNULL, cwd=thisdir)
+
+
 def _kill_docker_process():
     if _docker_process is not None:
         _docker_process.kill()

tests/test_session.py

@@ -31,6 +31,35 @@ async def test_admin(env):
         assert len(admin_users) == len(env.users.keys()), "Admin cannot see correct number of users"
         assert len(admin_sessions) == 2, "Admin cannot see correct number of oser sessions"
 
+async def test_auto_reconnect(env):
+    async with meshctrl.Session(env.mcurl, user="admin", password=env.users["admin"], ignore_ssl=True, auto_reconnect=True) as admin_session:
+        env.restart_mesh()
+        await asyncio.sleep(10)
+        await admin_session.ping(timeout=10)
+
+    # As above, but with proxy
+    async with meshctrl.Session("wss://" + env.dockerurl, user="admin", password=env.users["admin"], ignore_ssl=True, auto_reconnect=True, proxy=env.proxyurl) as admin_session:
+        env.restart_mesh()
+        for i in range(3):
+            try:
+                await admin_session.ping(timeout=10)
+            except:
+                continue
+            break
+        else:
+            raise Exception("Failed to reconnect")
+
+        env.restart_proxy()
+        for i in range(3):
+            try:
+                await admin_session.ping(timeout=10)
+            except* Exception as e:
+                pass
+            else:
+                break
+        else:
+            raise Exception("Failed to reconnect")
+
 
 async def test_users(env):
     try: