diff --git a/docs/auth.rst b/docs/auth.rst index f0d5ce0..24067f2 100644 --- a/docs/auth.rst +++ b/docs/auth.rst @@ -78,27 +78,10 @@ using them easy. Fetching a Token and Creating a Client -------------------------------------- -``schwab-py`` provides an easy implementation of the client-side login flow in -the ``auth`` package. It uses a `selenium -`__ webdriver to open the Schwab -authentication URL, take your login credentials, catch the post-login redirect, -and fetch a reusable token. It returns a fully-configured :ref:`client`, ready -to send API calls. It also handles token refreshing, and writes updated tokens -to the token file. - -These functions are webdriver-agnostic, meaning you can use whatever -webdriver-supported browser you have available on your system. You can find -information about available webdriver on the `Selenium documentation -`__. - -.. autofunction:: schwab.auth.client_from_login_flow - .. _manual_login: -If for some reason you cannot open a web browser, such as when running in a -cloud environment, the following function will guide you through the process of -manually creating a token by copy-pasting relevant URLs. +This function will guide you through the process of logging in and creating a +token. .. autofunction:: schwab.auth.client_from_manual_flow @@ -107,16 +90,10 @@ the login flow again. .. autofunction:: schwab.auth.client_from_token_file -The following is a convenient wrapper around these two methods, calling each -when appropriate: - -.. autofunction:: schwab.auth.easy_client - If you don't want to create a client and just want to fetch a token, you can use the ``schwab-generate-token.py`` script that's installed with the library. This method is particularly useful if you want to create your token on one machine -and use it on another. The script will attempt to open a web browser and perform -the login flow. If it fails, it will fall back to the manual login flow: +and use it on another. .. code-block:: bash @@ -173,34 +150,6 @@ can also `join our Discord server `__ to ask questio -.. _missing_chromedriver: - -++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ -``WebDriverException: Message: 'chromedriver' executable needs to be in PATH`` -++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - -When creating a ``schwab-py`` token using a webrowser-based method like -:func:`~schwab.auth.client_from_login_flow` or :func:`~schwab.auth.easy_client`, -the library must control the browser using `selenium -`__. This is a Python library that -sends commands to the browser to perform operations like load pages, inject -synthetic clicks, enter text, and so on. The component which is used to send -these commands is called a *driver*. - -Drivers are generally not part of the standard web browser installation, meaning -you must install them manually. If you're seeing this or a similar message, you -probably haven't installed the appropriate webdriver. Drivers are -available for most of the common web browsers, including `Chrome -`__, `Firefox -`__, and `Safari -`__. -Make sure you've installed the driver *before* attempting to create a token -using ``schwab-py``. - - -.. _invalid_grant: - - ++++++++++++++++++++++ Token Parsing Failures ++++++++++++++++++++++ @@ -228,26 +177,3 @@ you're confident is valid, please `file a ticket `__. Just remember, **never share your token file, not even with** ``schwab-py`` **developers**. Sharing the token file is as dangerous as sharing your Schwab username and password. - - -++++++++++++++++++++++++++++++ -What If I Can't Use a Browser? -++++++++++++++++++++++++++++++ - -Launching a browser can be inconvenient in some situations, most notably in -containerized applications running on a cloud provider. ``schwab-py`` supports -two alternatives to creating tokens by opening a web browser. - -Firstly, the :ref:`manual login flow` flow allows you to go -through the login flow on a different machine than the one on which -``schwab-py`` is running. Instead of starting the web browser and automatically -opening the relevant URLs, this flow allows you to manually copy-paste around -the URLs. It's a little more cumbersome, but it has no dependency on selenium. - -Alterately, you can take advantage of the fact that token files are portable. -Once you create a token on one machine, such as one where you can open a web -browser, you can easily copy that token file to another machine, such as your -application in the cloud. However, make sure you don't use the same token on -two machines. It is recommended to delete the token created on the -browser-capable machine as soon as it is copied to its destination. - diff --git a/schwab/auth.py b/schwab/auth.py index b98c5ac..8d93704 100644 --- a/schwab/auth.py +++ b/schwab/auth.py @@ -7,7 +7,6 @@ import json import logging import os -import pickle import sys import time import warnings @@ -37,43 +36,50 @@ def load_token(): get_logger().info('Loading token from file %s', token_path) with open(token_path, 'rb') as f: - token_data = f.read() - try: - return json.loads(token_data.decode()) - except ValueError: - get_logger().warning( - 'Unable to load JSON token from file %s, ' + - 'falling back to pickle', token_path) - return pickle.loads(token_data) + return json.load(f) return load_token -def client_from_token_file(token_path, api_key, app_secret, asyncio=False, - enforce_enums=True): +class TokenMetadata: ''' - Returns a session from an existing token file. The session will perform - an auth refresh as needed. It will also update the token on disk whenever - appropriate. - - :param token_path: Path to an existing token. Updated tokens will be written - to this path. If you do not yet have a token, use - :func:`~schwab.auth.client_from_login_flow` or - :func:`~schwab.auth.easy_client` to create one. - :param api_key: Your Schwab application's app key. - :param asyncio: If set to ``True``, this will enable async support allowing - the client to be used in an async environment. Defaults to - ``False`` - :param enforce_enums: Set it to ``False`` to disable the enum checks on ALL - the client methods. Only do it if you know you really - need it. For most users, it is advised to use enums - to avoid errors. + Provides the functionality required to maintain and update our view of the + token's metadata. ''' + def __init__(self, token, unwrapped_token_write_func): + # The token write function is ultimately stored in the session. When we + # get a new token we immediately wrap it in a new sesssion. We hold on + # to the unwrapped token writer function to allow us to inject the + # appropriate write function. + self.unwrapped_token_write_func = unwrapped_token_write_func - load = __token_loader(token_path) + # The current token. Updated whenever the wrapped token update function + # is called. + self.token = token - return client_from_access_functions( - api_key, app_secret, load, __update_token(token_path), asyncio=asyncio, - enforce_enums=enforce_enums) + @classmethod + def from_loaded_token(cls, token, unwrapped_token_write_func): + ''' + Returns a new ``TokenMetadata`` object extracted from the metadata of + the loaded token object. If the token has a legacy format which contains + no metadata, assign default values. + ''' + return TokenMetadata(token, unwrapped_token_write_func) + + def wrapped_token_write_func(self): + ''' + Returns a version of the unwrapped write function which wraps the token + in metadata and updates our view on the most recent token. + ''' + def wrapped_token_write_func(token, *args, **kwargs): + # If the write function is going to raise an exception, let it do so + # here before we update our reference to the current token. + ret = self.unwrapped_token_write_func(token, *args, **kwargs) + + self.token = token + + return ret + + return wrapped_token_write_func def __fetch_and_register_token_from_redirect( @@ -91,7 +97,7 @@ def __fetch_and_register_token_from_redirect( update_token = ( __update_token(token_path) if token_write_func is None else token_write_func) - metadata_manager = TokenMetadata(token, int(time.time()), update_token) + metadata_manager = TokenMetadata(token, update_token) update_token = metadata_manager.wrapped_token_write_func() update_token(token) @@ -120,94 +126,18 @@ async def oauth_client_update_token(t, *args, **kwargs): token_metadata=metadata_manager, enforce_enums=enforce_enums) -class RedirectTimeoutError(Exception): - pass - - -class TokenMetadata: - ''' - Provides the functionality required to maintain and update our view of the - token's metadata. - ''' - def __init__( - self, token, creation_timestamp, unwrapped_token_write_func=None): - self.creation_timestamp = creation_timestamp - - # The token write function is ultimately stored in the session. When we - # get a new token we immediately wrap it in a new sesssion. We hold on - # to the unwrapped token writer function to allow us to inject the - # appropriate write function. - self.unwrapped_token_write_func = unwrapped_token_write_func - - # The current token. Updated whenever the wrapped token update function - # is called. - self.token = token - - @classmethod - def from_loaded_token(cls, token, app_secret, unwrapped_token_write_func=None): - ''' - Returns a new ``TokenMetadata`` object extracted from the metadata of - the loaded token object. If the token has a legacy format which contains - no metadata, assign default values. - ''' - logger = get_logger() - logger.info( - 'Loaded metadata aware token with creation timestamp %s', - token['creation_timestamp']) - return TokenMetadata( - token['token'], - token['creation_timestamp'], - unwrapped_token_write_func) - - def wrapped_token_write_func(self): - ''' - Returns a version of the unwrapped write function which wraps the token - in metadata and updates our view on the most recent token. - ''' - def wrapped_token_write_func(token, *args, **kwargs): - # If the write function is going to raise an exception, let it do so - # here before we update our reference to the current token. - ret = self.unwrapped_token_write_func( - self.wrap_token_in_metadata(token), *args, **kwargs) - - self.token = token - - return ret - - return wrapped_token_write_func - - def wrap_token_in_metadata(self, token): - return { - 'creation_timestamp': self.creation_timestamp, - 'token': token, - } - - -# TODO: Raise an exception when passing both token_path and token_write_func -def client_from_login_flow(webdriver, api_key, app_secret, callback_url, token_path, - redirect_wait_time_seconds=0.1, max_waits=3000, - asyncio=False, token_write_func=None, +def client_from_token_file(token_path, api_key, app_secret, asyncio=False, enforce_enums=True): ''' - Uses the webdriver to perform an OAuth webapp login flow and creates a - client wrapped around the resulting token. The client will be configured to - refresh the token as necessary, writing each updated version to - ``token_path``. - - **Warning:** Schwab appears to block logins performed within a webdriver. - This library has been included as a direct copy from ``tda-api``, but it may - be removed in the future. + Returns a session from an existing token file. The session will perform + an auth refresh as needed. It will also update the token on disk whenever + appropriate. - :param webdriver: `selenium `__ - webdriver which will be used to perform the login flow. + :param token_path: Path to an existing token. Updated tokens will be written + to this path. If you do not yet have a token, use + :func:`~schwab.auth.client_from_login_flow` or + :func:`~schwab.auth.easy_client` to create one. :param api_key: Your Schwab application's app key. - :param callback_url: Your Schwab application's callback URL. Note this must - *exactly* match the value you've entered in your - application configuration, otherwise login will fail - with a security error. - :param token_path: Path to which the new token will be written. If the token - file already exists, it will be overwritten with a new - one. Updated tokens will be written to this path as well. :param asyncio: If set to ``True``, this will enable async support allowing the client to be used in an async environment. Defaults to ``False`` @@ -216,51 +146,12 @@ def client_from_login_flow(webdriver, api_key, app_secret, callback_url, token_p need it. For most users, it is advised to use enums to avoid errors. ''' - get_logger().info('Creating new token with redirect URL \'%s\' ' + - 'and token path \'%s\'', callback_url, token_path) - - oauth = OAuth2Client(api_key, callback_url=callback_url) - authorization_url, state = oauth.create_authorization_url( - 'https://api.schwabapi.com/v1/oauth/authorize') - - # Open the login page and wait for the redirect - print('\n**************************************************************\n') - print('Opening the login page in a webdriver. Please use this window to', - 'log in. Successful login will be detected automatically.') - print() - print('If you encounter any issues, see here for troubleshooting: ' + - 'https://schwab-py.readthedocs.io/en/stable/auth.html' + - '#troubleshooting') - print('\n**************************************************************\n') - - webdriver.get(authorization_url) - - # Tolerate redirects to HTTPS on the callback URL - if callback_url.startswith('http://'): - print(('WARNING: Your callback URL ({}) will transmit data over HTTP, ' + - 'which is a potentially severe security vulnerability. ' + - 'Please go to your app\'s configuration with Schwab ' + - 'and update your callback URL to begin with \'https\' ' + - 'to stop seeing this message.').format(callback_url)) - callback_urls = (callback_url, 'https' + callback_url[4:]) - else: - callback_urls = (callback_url,) - - # Wait until the current URL starts with the callback URL - current_url = '' - num_waits = 0 - while not any(current_url.startswith(r_url) for r_url in callback_urls): - current_url = webdriver.current_url - - if num_waits > max_waits: - raise RedirectTimeoutError('timed out waiting for redirect') - time.sleep(redirect_wait_time_seconds) - num_waits += 1 + load = __token_loader(token_path) - return __fetch_and_register_token_from_redirect( - oauth, current_url, api_key, app_secret, token_path, token_write_func, - asyncio, enforce_enums=enforce_enums) + return client_from_access_functions( + api_key, app_secret, load, __update_token(token_path), asyncio=asyncio, + enforce_enums=enforce_enums) def client_from_manual_flow(api_key, app_secret, callback_url, token_path, @@ -336,68 +227,6 @@ def client_from_manual_flow(api_key, app_secret, callback_url, token_path, asyncio, enforce_enums=enforce_enums) -def easy_client(api_key, app_secret, callback_url, token_path, - webdriver_func=None, asyncio=False, enforce_enums=True): - '''Convenient wrapper around :func:`client_from_login_flow` and - :func:`client_from_token_file`. If ``token_path`` exists, loads the token - from it. Otherwise open a login flow to fetch a new token. Returns a client - configured to refresh the token to ``token_path``. - - **Warning:** Schwab appears to block logins performed within a webdriver. - This library has been included as a direct copy from ``tda-api``, but it may - be removed in the future. - - *Reminder:* You should never create the token file yourself or modify it in - any way. If ``token_path`` refers to an existing file, this method will - assume that file is valid token and will attempt to parse it. - - :param api_key: Your Schwab application's app key. - :param callback_url: Your Schwab application's redirect URL. Note this must - *exactly* match the value you've entered in your - application configuration, otherwise login will fail - with a security error. - :param token_path: Path that new token will be read from and written to. If - If this file exists, this method will assume it's valid - and will attempt to parse it as a token. If it does not, - this method will create a new one using - :func:`~schwab.auth.client_from_login_flow`. Updated tokens - will be written to this path as well. - :param webdriver_func: Function that returns a webdriver for use in fetching - a new token. Will only be called if the token file - cannot be found. - :param asyncio: If set to ``True``, this will enable async support allowing - the client to be used in an async environment. Defaults to - ``False`` - :param enforce_enums: Set it to ``False`` to disable the enum checks on ALL - the client methods. Only do it if you know you really - need it. For most users, it is advised to use enums - to avoid errors. - ''' - logger = get_logger() - - if os.path.isfile(token_path): - c = client_from_token_file(token_path, api_key, app_secret, - asyncio=asyncio, enforce_enums=enforce_enums) - logger.info( - 'Returning client loaded from token file \'%s\'', token_path) - return c - else: - logger.warning('Failed to find token file \'%s\'', token_path) - - if webdriver_func is not None: - with webdriver_func() as driver: - c = client_from_login_flow( - driver, api_key, app_secret, callback_url, token_path, - asyncio=asyncio, enforce_enums=enforce_enums) - logger.info( - 'Returning client fetched using webdriver, writing' + - 'token to \'%s\'', token_path) - return c - else: - logger.error('No webdriver_func set, cannot fetch token') - sys.exit(1) - - def client_from_access_functions(api_key, app_secret, token_read_func, token_write_func, asyncio=False, enforce_enums=True): @@ -412,7 +241,7 @@ def client_from_access_functions(api_key, app_secret, token_read_func, Users are free to customize how they represent the token file. In theory, since they have direct access to the token, they can get creative about how they store it and fetch it. In practice, it is *highly* recommended to - simply accept the token object and use ``pickle`` to serialize and + simply accept the token object and use ``json`` to serialize and deserialize it, without inspecting it in any way. Note the read and write methods must take particular arguments. Please see @@ -437,16 +266,11 @@ def client_from_access_functions(api_key, app_secret, token_read_func, token = token_read_func() # Extract metadata and unpack the token, if necessary - metadata = TokenMetadata.from_loaded_token( - token, app_secret, token_write_func) - token = token['token'] + metadata = TokenMetadata.from_loaded_token(token, token_write_func) # Don't emit token details in debug logs register_redactions(token) - # Return a new session configured to refresh credentials - #api_key = _normalize_api_key(api_key) - wrapped_token_write_func = metadata.wrapped_token_write_func() if asyncio: diff --git a/tests/auth_test.py b/tests/auth_test.py new file mode 100644 index 0000000..d71c72c --- /dev/null +++ b/tests/auth_test.py @@ -0,0 +1,420 @@ +from schwab import auth +from .utils import ( + AnyStringWith, + MockAsyncOAuthClient, + MockOAuthClient, + no_duplicates +) +from unittest.mock import patch, ANY, MagicMock +from unittest.mock import ANY as _ + +import json +import os +import tempfile +import unittest + + +API_KEY = 'APIKEY' +APP_SECRET = '0x5EC07' +MOCK_NOW = 1613745082 +REDIRECT_URL = 'https://redirect.url.com' + + +class ClientFromTokenFileTest(unittest.TestCase): + + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory() + self.json_path = os.path.join(self.tmp_dir.name, 'token.json') + self.token = {'token': 'yes'} + + def write_token(self): + with open(self.json_path, 'w') as f: + json.dump(self.token, f) + + @no_duplicates + def test_no_such_file(self): + with self.assertRaises(FileNotFoundError): + auth.client_from_token_file(self.json_path, API_KEY, APP_SECRET) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_json_loads(self, async_session, sync_session, client): + self.write_token() + + client.return_value = 'returned client' + + self.assertEqual('returned client', + auth.client_from_token_file( + self.json_path, API_KEY, APP_SECRET)) + client.assert_called_once_with(API_KEY, _, token_metadata=_, + enforce_enums=_) + sync_session.assert_called_once_with( + API_KEY, + client_secret=APP_SECRET, + token=self.token, + token_endpoint=_, + update_token=_) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_update_token_updates_token( + self, async_session, sync_session, client): + self.write_token() + + auth.client_from_token_file(self.json_path, API_KEY, APP_SECRET) + sync_session.assert_called_once() + + session_call = sync_session.mock_calls[0] + update_token = session_call[2]['update_token'] + + updated_token = {'updated': 'token'} + update_token(updated_token) + with open(self.json_path, 'r') as f: + self.assertEqual(json.load(f), updated_token) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_enforce_enums_being_disabled(self, async_session, sync_session, client): + self.write_token() + + client.return_value = 'returned client' + + self.assertEqual('returned client', + auth.client_from_token_file( + self.json_path, API_KEY, APP_SECRET, + enforce_enums=False)) + client.assert_called_once_with(API_KEY, _, token_metadata=_, + enforce_enums=False) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_enforce_enums_being_enabled(self, async_session, sync_session, client): + self.write_token() + + client.return_value = 'returned client' + + self.assertEqual('returned client', + auth.client_from_token_file( + self.json_path, API_KEY, APP_SECRET)) + client.assert_called_once_with(API_KEY, _, token_metadata=_, + enforce_enums=True) + + +class ClientFromAccessFunctionsTest(unittest.TestCase): + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_success_with_write_func( + self, async_session, sync_session, client): + token = {'token': 'yes'} + + token_read_func = MagicMock() + token_read_func.return_value = token + + token_writes = [] + + def token_write_func(token): + token_writes.append(token) + + client.return_value = 'returned client' + self.assertEqual('returned client', + auth.client_from_access_functions( + API_KEY, + APP_SECRET, + token_read_func, + token_write_func)) + + sync_session.assert_called_once_with( + API_KEY, + client_secret=APP_SECRET, + token=token, + token_endpoint=_, + update_token=_) + token_read_func.assert_called_once() + + # Verify that the write function is called when the updater is called + session_call = sync_session.mock_calls[0] + update_token = session_call[2]['update_token'] + + + update_token(token) + self.assertEqual([token], token_writes) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_success_with_write_func_metadata_aware_token( + self, async_session, sync_session, client): + token = {'token': 'yes'} + + token_read_func = MagicMock() + token_read_func.return_value = token + + token_writes = [] + + def token_write_func(token): + token_writes.append(token) + + client.return_value = 'returned client' + self.assertEqual('returned client', + auth.client_from_access_functions( + API_KEY, + APP_SECRET, + token_read_func, + token_write_func)) + + sync_session.assert_called_once_with( + API_KEY, + client_secret=APP_SECRET, + token=token, + token_endpoint=_, + update_token=_) + token_read_func.assert_called_once() + + # Verify that the write function is called when the updater is called + session_call = sync_session.mock_calls[0] + update_token = session_call[2]['update_token'] + + update_token(token) + self.assertEqual([token], token_writes) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_success_with_enforce_enums_disabled( + self, async_session, sync_session, client): + token = {'token': 'yes'} + + token_read_func = MagicMock() + token_read_func.return_value = token + + token_writes = [] + + def token_write_func(token): + token_writes.append(token) + + client.return_value = 'returned client' + self.assertEqual('returned client', + auth.client_from_access_functions( + API_KEY, + APP_SECRET, + token_read_func, + token_write_func, enforce_enums=False)) + + client.assert_called_once_with( + API_KEY, _, token_metadata=_, enforce_enums=False) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + def test_success_with_enforce_enums_enabled( + self, async_session, sync_session, client): + token = {'token': 'yes'} + + token_read_func = MagicMock() + token_read_func.return_value = token + + token_writes = [] + + def token_write_func(token): + token_writes.append(token) + + client.return_value = 'returned client' + self.assertEqual('returned client', + auth.client_from_access_functions( + API_KEY, + APP_SECRET, + token_read_func, + token_write_func)) + + client.assert_called_once_with( + API_KEY, _, token_metadata=_, enforce_enums=True) + + + +class ClientFromManualFlow(unittest.TestCase): + + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory() + self.json_path = os.path.join(self.tmp_dir.name, 'token.json') + self.token = {'token': 'yes'} + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + @patch('schwab.auth.prompt') + @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW)) + def test_no_token_file( + self, prompt_func, 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.token + + client.return_value = 'returned client' + prompt_func.return_value = 'http://redirect.url.com/?data' + + self.assertEqual('returned client', + auth.client_from_manual_flow( + API_KEY, APP_SECRET, REDIRECT_URL, self.json_path)) + + with open(self.json_path, 'r') as f: + self.assertEqual(self.token, json.load(f)) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + @patch('schwab.auth.prompt') + @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW)) + def test_custom_token_write_func( + self, prompt_func, 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.token + + webdriver = MagicMock() + webdriver.current_url = REDIRECT_URL + '/token_params' + + client.return_value = 'returned client' + prompt_func.return_value = 'http://redirect.url.com/?data' + + token_writes = [] + + def dummy_token_write_func(token): + token_writes.append(token) + + self.assertEqual('returned client', + auth.client_from_manual_flow( + API_KEY, APP_SECRET, REDIRECT_URL, + self.json_path, + token_write_func=dummy_token_write_func)) + + sync_session.assert_called_with( + _, client_secret=APP_SECRET, token=_, update_token=_) + + self.assertEqual([self.token], token_writes) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + @patch('schwab.auth.prompt') + @patch('builtins.print') + @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW)) + def test_print_warning_on_http_redirect_uri( + self, print_func, prompt_func, async_session, sync_session, client): + AUTH_URL = 'https://auth.url.com' + + redirect_url = 'http://redirect.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.token + + client.return_value = 'returned client' + prompt_func.return_value = 'http://redirect.url.com/?data' + + self.assertEqual('returned client', + auth.client_from_manual_flow( + API_KEY, APP_SECRET, redirect_url, self.json_path)) + + with open(self.json_path, 'r') as f: + self.assertEqual(self.token, json.load(f)) + + print_func.assert_any_call(AnyStringWith('will transmit data over HTTP')) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + @patch('schwab.auth.prompt') + @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW)) + def test_enforce_enums_disabled( + self, prompt_func, 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.token + + client.return_value = 'returned client' + prompt_func.return_value = 'http://redirect.url.com/?data' + + self.assertEqual('returned client', + auth.client_from_manual_flow( + API_KEY, APP_SECRET, REDIRECT_URL, self.json_path, + enforce_enums=False)) + + client.assert_called_once_with(API_KEY, _, token_metadata=_, + enforce_enums=False) + + @no_duplicates + @patch('schwab.auth.Client') + @patch('schwab.auth.OAuth2Client', new_callable=MockOAuthClient) + @patch('schwab.auth.AsyncOAuth2Client', new_callable=MockAsyncOAuthClient) + @patch('schwab.auth.prompt') + @patch('time.time', unittest.mock.MagicMock(return_value=MOCK_NOW)) + def test_enforce_enums_enabled( + self, prompt_func, 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.token + + client.return_value = 'returned client' + prompt_func.return_value = 'http://redirect.url.com/?data' + + self.assertEqual('returned client', + auth.client_from_manual_flow( + API_KEY, APP_SECRET, REDIRECT_URL, self.json_path)) + + client.assert_called_once_with(API_KEY, _, token_metadata=_, + enforce_enums=True) + + +class TokenMetadataTest(unittest.TestCase): + + @no_duplicates + def test_from_loaded_token(self): + token = {'token': 'yes'} + + metadata = auth.TokenMetadata.from_loaded_token( + token, unwrapped_token_write_func=None) + self.assertEqual(metadata.token, token) + + + @no_duplicates + def test_wrapped_token_write_func_updates_stored_token(self): + token = {'token': 'yes'} + + updated = [False] + def update_token(token): + updated[0] = True + + metadata = auth.TokenMetadata.from_loaded_token( + token, unwrapped_token_write_func=update_token) + + new_token = {'updated': 'yes'} + metadata.wrapped_token_write_func()(new_token) + + self.assertTrue(updated[0]) + self.assertEqual(new_token, metadata.token)