From d96a693d7ea2d36055ffa6b0cc4a8bf34a4a3ffc Mon Sep 17 00:00:00 2001
From: Marcel van der Veldt <m.vanderveldt@outlook.com>
Date: Wed, 2 Sep 2020 23:05:46 +0200
Subject: [PATCH] few fixes

---
 hass_client/__init__.py | 22 +++++++++++++++++++---
 setup.py                |  8 +++-----
 test.py                 |  5 ++++-
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/hass_client/__init__.py b/hass_client/__init__.py
index 8e0b86d..c07da55 100644
--- a/hass_client/__init__.py
+++ b/hass_client/__init__.py
@@ -27,7 +27,6 @@ def __init__(self, url: str = None, token: str = None):
             Initialize the connection to HomeAssistant.
                 :param url: full url to the HomeAssistant instance.
                 :param token: a long lived token.
-            Note: Leave url and token empty for auto detection on supervisor addon's.
         """
         self._loop = asyncio.get_event_loop()
         self._states = {}
@@ -48,8 +47,17 @@ def __init__(self, url: str = None, token: str = None):
         self._initial_state_received = False
         self._connected_callback = None
         self._event_listeners = []
+        self._ws_task = None
 
-    async def connect(self):
+    def connect(self):
+        """Start the connection."""
+        self._loop.create_task(self.async_connect())
+
+    def close(self):
+        """Close the connection."""
+        self._loop.create_task(self.async_close())
+
+    async def async_connect(self):
         """Start the connection."""
         if not self._loop.is_running:
             raise RuntimeError("A running eventloop is required!")
@@ -58,7 +66,15 @@ async def connect(self):
         self._http_session = aiohttp.ClientSession(
             loop=self._loop, connector=aiohttp.TCPConnector()
         )
-        self._loop.create_task(self.__async_hass_websocket())
+        self._ws_task = self._loop.create_task(self.__async_hass_websocket())
+
+    async def async_close(self):
+        """Close/stop the connection."""
+        if self._ws_task:
+            self._ws_task.cancel()
+        if self._http_session:
+            await self._http_session.close()
+        LOGGER.info("Disconnected from Home Assistant")
 
     def register_event_callback(
         self,
diff --git a/setup.py b/setup.py
index af82eb5..2056dfb 100644
--- a/setup.py
+++ b/setup.py
@@ -5,12 +5,10 @@
 
 PROJECT_DIR = Path(__file__).parent.resolve()
 README_FILE = PROJECT_DIR / "README.md"
-VERSION = "0.0.1"
+VERSION = "0.0.2"
 
 with open("requirements.txt") as f:
     INSTALL_REQUIRES = f.read().splitlines()
-if os.name != "nt":
-    INSTALL_REQUIRES.append("uvloop")
 
 PACKAGE_FILES = []
 for (path, directories, filenames) in os.walk('hass_client/'):
@@ -20,8 +18,8 @@
 setup(
     name="hass_client",
     version=VERSION,
-    url="https://github.com/marcelveldt/hass-client",
-    download_url="https://github.com/marcelveldt/hass-client",
+    url="https://github.com/marcelveldt/python-hass-client",
+    download_url="https://github.com/marcelveldt/python-hass-client",
     author="Marcel van der Veldt",
     author_email="m.vanderveldt@outlook.com",
     description="Basic client for connecting to Home Assistant over websockets and REST.",
diff --git a/test.py b/test.py
index 1a1f92e..d9dcbcd 100644
--- a/test.py
+++ b/test.py
@@ -33,7 +33,10 @@ async def hass_event(event, event_details):
     hass.register_event_callback(hass_event, )
 
     async def run():
-        await hass.connect()
+        await hass.async_connect()
+        await asyncio.sleep(10)
+        await hass.async_close()
+        loop.stop()
 
     try:
         loop.create_task(run())