From 98f7f8fbcf6d0655db848ad8a4e0a4ccca629d17 Mon Sep 17 00:00:00 2001
From: Jason Lawrence <jjlawren@users.noreply.github.com>
Date: Fri, 27 Sep 2019 15:56:06 -0500
Subject: [PATCH] Allow session to be passed in, add async context manager

---
 README.md   |  8 ++++----
 plexauth.py | 45 +++++++++++++++++++++++++++++++--------------
 setup.py    |  2 +-
 3 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index ff44302..d5fb3d1 100644
--- a/README.md
+++ b/README.md
@@ -18,10 +18,10 @@ PAYLOAD = {
 }
 
 async def main():
-    plexauth = PlexAuth(PAYLOAD)
-    await plexauth.initiate_auth()
-    print("Complete auth at URL: {}".format(plexauth.auth_url()))
-    token = await plexauth.token()
+    async with PlexAuth(PAYLOAD) as plexauth:
+        await plexauth.initiate_auth()
+        print("Complete auth at URL: {}".format(plexauth.auth_url()))
+        token = await plexauth.token()
     
     if token:
         print("Token: {}".format(token))
diff --git a/plexauth.py b/plexauth.py
index 90083ca..8d722d6 100755
--- a/plexauth.py
+++ b/plexauth.py
@@ -5,7 +5,7 @@
 import urllib.parse
 import uuid
 
-__version__ = '0.0.1'
+__version__ = '0.0.3'
 
 CODES_URL = 'https://plex.tv/api/v2/pins.json?strong=true'
 AUTH_URL = 'https://app.plex.tv/auth#!?{}'
@@ -13,7 +13,7 @@
 
 class PlexAuth():
 
-    def __init__(self, payload):
+    def __init__(self, payload, session=None):
         '''Create PlexAuth instance.'''
         self._client_identifier = str(uuid.uuid4())
         self._code = None
@@ -21,13 +21,18 @@ def __init__(self, payload):
         self._payload = payload
         self._payload['X-Plex-Client-Identifier'] = self._client_identifier
 
+        self._local_session = False
+        self._session = session
+        if session is None:
+            self._session = aiohttp.ClientSession()
+            self._local_session = True
+
     async def initiate_auth(self):
         '''Request codes needed to create an auth URL. Starts external timeout.'''
-        async with aiohttp.ClientSession() as session:
-            async with session.post(CODES_URL, data=self._payload) as resp:
-                response = await resp.json()
-                self._code = response['code']
-                self._identifier = response['id']
+        async with self._session.post(CODES_URL, data=self._payload) as resp:
+            response = await resp.json()
+            self._code = response['code']
+            self._identifier = response['id']
 
     def auth_url(self, forward_url=None):
         '''Return an auth URL for the user to follow.'''
@@ -44,13 +49,12 @@ def auth_url(self, forward_url=None):
     async def request_auth_token(self):
         '''Request an auth token from Plex.'''
         token_url = TOKEN_URL.format(self._code)
-        async with aiohttp.ClientSession() as session:
-            payload = dict(self._payload)
-            payload['Accept'] = 'application/json'
-            async with session.get(TOKEN_URL.format(self._identifier), headers=payload) as resp:
-                response = await resp.json()
-                token = response['authToken']
-                return token
+        payload = dict(self._payload)
+        payload['Accept'] = 'application/json'
+        async with self._session.get(TOKEN_URL.format(self._identifier), headers=payload) as resp:
+            response = await resp.json()
+            token = response['authToken']
+            return token
 
     async def token(self, timeout=60):
         '''Poll Plex endpoint until a token is retrieved or times out.'''
@@ -64,3 +68,16 @@ async def token(self, timeout=60):
                 break_loop = True
 
         return token
+
+    async def close(self):
+        """Close open client session."""
+        if self._local_session:
+            await self._session.close()
+
+    async def __aenter__(self):
+        """Async enter."""
+        return self
+
+    async def __aexit__(self, *exc_info):
+        """Async exit."""
+        await self.close()
diff --git a/setup.py b/setup.py
index ee1985e..85ad59e 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 with open('README.md') as f:
     long_description = f.read()
 
-VERSION="0.0.1"
+VERSION="0.0.3"
 
 setup(
     name='plexauth',