Adds message and interactivity

schwab/auth.py

@@ -1,9 +1,7 @@
-
-
 from authlib.integrations.httpx_client import AsyncOAuth2Client, OAuth2Client
 from prompt_toolkit import prompt
 
-import urllib
+import atexit
 import json
 import logging
 import multiprocessing
@@ -13,6 +11,7 @@
 import requests
 import sys
 import time
+import urllib
 import urllib3
 import warnings
 import webbrowser
@@ -184,7 +183,16 @@ def handle_token():
     def status():
         return 'running'
 
-    app.run(port=callback_port, ssl_context='adhoc')
+    # Wrap this call in some hackery to suppress the flask startup messages
+    with open(os.devnull, 'w') as devnull:
+        import logging
+        log = logging.getLogger('werkzeug')
+        log.setLevel(logging.ERROR)
+
+        old_stdout = sys.stdout
+        sys.stdout = devnull
+        app.run(port=callback_port, ssl_context='adhoc')
+        sys.stdout = old_stdout
 
 
 class RedirectTimeoutError(Exception):
@@ -193,10 +201,14 @@ class RedirectTimeoutError(Exception):
 class RedirectServerExitedError(Exception):
     pass
 
+# Capture the real time.time so that we can use it in server initialization 
+# while simultaneously mocking it in testing
+__TIME_TIME = time.time
 
 def client_from_login_flow(api_key, app_secret, callback_url, token_path,
                            asyncio=False, enforce_enums=False, 
-                           token_write_func=None, callback_timeout=300.0):
+                           token_write_func=None, callback_timeout=300.0,
+                           interactive=True):
     # TODO: documentation
 
     # Start the server
@@ -217,10 +229,13 @@ def client_from_login_flow(api_key, app_secret, callback_url, token_path,
             target=__run_client_from_login_flow_server,
             args=(output_queue, callback_port, callback_path))
 
-    print('Running a server to intercept the callback. Please ignore the ' +
-          'following debug messages:')
-    print()
     server.start()
+    def kill_server():
+        try:
+            psutil.Process(server.pid).kill()
+        except psutil.NoSuchProcess:
+            pass
+    atexit.register(kill_server)
 
     # Wait until the server successfully starts
     while True:
@@ -252,10 +267,37 @@ def client_from_login_flow(api_key, app_secret, callback_url, token_path,
     authorization_url, state = oauth.create_authorization_url(
         'https://api.schwabapi.com/v1/oauth/authorize')
 
+    if interactive:
+        print()
+        print('**************************************************************')
+        print()
+        print('This is the browser-assisted login and token creation flow for')
+        print('schwab-py. This flow automatically opens the login page on your')
+        print('browser, captures the resulting OAuth callback, and creates a token')
+        print('using the result.')
+        print()
+        print('IMPORTANT: Your browser will give you a security warning about an')
+        print('invalid certificate prior to issuing the redirect. This is because')
+        print('schwab-py has started a server on your machine to receive the OAuth')
+        print('redirect using a self-signed SSL certificate. You can ignore that')
+        print('warning, but make sure to first check that the URL matches your')
+        print('callback URL. As a reminder, your callback URL is:')
+        print()
+        print('>>',callback_url)
+        print()
+        print('See here to learn more: TODO<add a documentation URL>')
+        print()
+        print('If you encounter any issues, see here for troubleshooting:')
+        print('https://schwab-py.readthedocs.io/en/latest/auth.html#troubleshooting')
+        print('\n**************************************************************')
+        print()
+        prompt('Press ENTER to open the browser. Note you can run ' +
+              'client_from_login_flow with interactive=False to skip this input')
+
     webbrowser.open(authorization_url)
 
     # Wait for a response
-    now = time.time()
+    now = __TIME_TIME()
     timeout_time = now + callback_timeout
     callback_url = None
     while now < timeout_time:
@@ -267,7 +309,7 @@ def client_from_login_flow(api_key, app_secret, callback_url, token_path,
         except queue.Empty:
             pass
 
-        now = time.time()
+        now = __TIME_TIME()
 
     # Clean up and create the client
     psutil.Process(server.pid).kill()

tests/auth_test.py

@@ -39,6 +39,7 @@ def setUp(self):
     @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient)
     @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient)
     @patch('schwab.auth.webbrowser.open', new_callable=MagicMock)
+    @patch('schwab.auth.prompt', unittest.mock.MagicMock(return_value=''))
     @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW))
     def test_create_token_file(
             self, mock_webbrowser_open, async_session, sync_session, client):
@@ -70,6 +71,43 @@ def test_create_token_file(
     @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient)
     @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient)
     @patch('schwab.auth.webbrowser.open', new_callable=MagicMock)
+    @patch('schwab.auth.prompt')
+    @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW))
+    def test_create_token_file_not_interactive(
+            self, mock_prompt,mock_webbrowser_open, async_session, sync_session,
+            client):
+        AUTH_URL = 'https://auth.url.com'
+
+        sync_session.return_value = sync_session
+        sync_session.create_authorization_url.return_value = AUTH_URL, None
+        sync_session.fetch_token.return_value = self.raw_token
+
+        callback_url = 'https://127.0.0.1:6969/callback'
+
+        mock_webbrowser_open.side_effect = \
+                lambda auth_url: requests.get(
+                        'https://127.0.0.1:6969/callback', verify=False)
+
+        client.return_value = 'returned client'
+
+        auth.client_from_login_flow(
+                API_KEY, APP_SECRET, callback_url, self.token_path, 
+                interactive=False)
+
+        with open(self.token_path, 'r') as f:
+            self.assertEqual({
+                'creation_timestamp': MOCK_NOW,
+                'token': self.raw_token
+            }, json.load(f))
+
+        mock_prompt.assert_not_called()
+
+
+    @patch('schwab.auth.Client')
+    @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient)
+    @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient)
+    @patch('schwab.auth.webbrowser.open', new_callable=MagicMock)
+    @patch('schwab.auth.prompt', unittest.mock.MagicMock(return_value=''))
     @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW))
     def test_create_token_file_root_callback_url(
             self, mock_webbrowser_open, async_session, sync_session, client):
@@ -143,6 +181,7 @@ def test_unprivileged_start_on_port_80(
     @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient)
     @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient)
     @patch('schwab.auth.webbrowser.open', new_callable=MagicMock)
+    @patch('schwab.auth.prompt', unittest.mock.MagicMock(return_value=''))
     def test_time_out_waiting_for_request(
             self, mock_webbrowser_open, async_session, sync_session, client):
         AUTH_URL = 'https://auth.url.com'