From 8957855c81ff423f03946bff457429ce23eb5409 Mon Sep 17 00:00:00 2001
From: renemarc <renemarc@gmail.com>
Date: Thu, 26 Mar 2020 14:39:23 -0400
Subject: [PATCH] Fix session close handling in command-line client

---
 pyebox/__main__.py |  3 ---
 pyebox/client.py   | 11 ++++++-----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/pyebox/__main__.py b/pyebox/__main__.py
index a1dc912..f67c193 100644
--- a/pyebox/__main__.py
+++ b/pyebox/__main__.py
@@ -61,16 +61,13 @@ def main():
         loop.run_until_complete(task)
     except PyEboxError as exp:
         print(exp)
-        client.close_session()
         return
     if not client.get_data():
-        client.close_session()
         return
     if args.json:
         print(json.dumps(client.get_data()))
     else:
         _format_output(args.username, client.get_data())
-    client.close_session()
 
 
 if __name__ == '__main__':
diff --git a/pyebox/client.py b/pyebox/client.py
index 349c2a2..ecb9b57 100644
--- a/pyebox/client.py
+++ b/pyebox/client.py
@@ -1,8 +1,6 @@
 """
 Pyebox
 """
-import asyncio
-import json
 import logging
 import re
 import async_timeout
@@ -38,7 +36,8 @@ class PyEboxError(Exception):
 
 class EboxClient(object):
 
-    def __init__(self, username, password, timeout=REQUESTS_TIMEOUT, session=None):
+    def __init__(self, username, password, timeout=REQUESTS_TIMEOUT,
+                 session=None):
         """Initialize the client object."""
         self.username = username
         self.password = password
@@ -184,14 +183,16 @@ async def fetch_data(self):
         # Merge data
         self._data.update(home_data)
         self._data.update(usage_data)
+        # Close http session
+        await self.close_session()
 
     def get_data(self):
         """Return collected data"""
         return self._data
 
-    def close_session(self):
+    async def close_session(self):
         """Close current session."""
         if not self._session.closed:
             if self._session._connector_owner:
-                self._session._connector.close()
+                await self._session._connector.close()
             self._session._connector = None