Skip to content

Commit

Permalink
updated script
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexgolec
alexgolec committed May 14, 2024
1 parent af0138f commit fce52ee
Show file tree
Hide file tree
Showing 2 changed files with 135 additions and 19 deletions.
57 changes: 49 additions & 8 deletions schwab/scripts/orders_codegen.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import argparse
import httpx
import json

from schwab.auth import client_from_token_file
Expand All @@ -13,25 +14,65 @@ def latest_order_main(sys_args):
required.add_argument(
'--token_file', required=True, help='Path to token file')
required.add_argument('--api_key', required=True)
required.add_argument('--app_secret', required=True)

parser.add_argument('--account_id', type=int,
account_spec_group = parser.add_mutually_exclusive_group()
account_spec_group.add_argument('--account_id', type=int,
help='Restrict lookups to a specific account ID')

account_spec_group.add_argument('--account_hash', type=str,
help='Restrict lookups to the account with the specified hash')

args = parser.parse_args(args=sys_args)
client = client_from_token_file(args.token_file, args.api_key)
client = client_from_token_file(
args.token_file, args.app_secret, args.api_key)

# If the account ID is specified, find the corresponding account hash
if args.account_id is not None:
r = client.get_account_numbers()
assert r.status_code == httpx.codes.OK

for val in r.json():
if val['accountNumber'] == str(args.account_id):
account_hash = val['hashValue']
break
else:
print(('Failed to find account has for account ID {}. Searched ' +

Check warning on line 40 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L40 

Added line #L40 was not covered by tests
'the following accounts:\n{}').format(
args.account_id, json.dumps(r.json(), indent=4)))
return -1

Check warning on line 43 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L43 

Added line #L43 was not covered by tests
else:
account_hash = args.account_hash


# Fetch orders
def get_orders(method):
r = method()
if r.status_code != httpx.codes.OK:
print(('Returned HTTP status code {}. This is most often caused ' +

Check warning on line 52 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L52 

Added line #L52 was not covered by tests
'by an invalid account ID or hash.').format(r.status_code))
return None

Check warning on line 54 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L54 

Added line #L54 was not covered by tests
return r.json()

if account_hash is not None:
orders = get_orders(lambda: client.get_orders_for_account(account_hash))
if orders is None:
return -1

Check warning on line 60 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L60 

Added line #L60 was not covered by tests

if args.account_id:
orders = client.get_orders_by_path(args.account_id).json()
if 'error' in orders:
print(('TDA returned error: "{}", This is most often caused by ' +
'an invalid account ID').format(orders['error']))
print(('Schwab returned error: "{}", This is most often caused ' +
'by an invalid account ID or hash').format(orders['error']))
return -1
else:
orders = client.get_orders_by_query().json()
orders = get_orders(lambda: client.get_orders_for_all_linked_accounts())
if orders is None:
return -1

Check warning on line 69 in schwab/scripts/orders_codegen.py

View check run for this annotation

Codecov / codecov/patch

schwab/scripts/orders_codegen.py#L69 

Added line #L69 was not covered by tests

if 'error' in orders:
print('TDA returned error: "{}"'.format(orders['error']))
print('Schwab returned error: "{}"'.format(orders['error']))
return -1

# Construct and emit order code
if orders:
order = sorted(orders, key=lambda o: -o['orderId'])[0]
print('# Order ID', order['orderId'])
Expand Down
97 changes: 86 additions & 11 deletions tests/scripts/orders_codegen_test.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import httpx
import subprocess
import unittest
from unittest.mock import call, MagicMock, patch
Expand Down Expand Up @@ -31,6 +32,8 @@ def test_success_no_account_id(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')

orders = [
{'orderId': 201},
Expand All @@ -41,7 +44,8 @@ def test_success_no_account_id(

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_query.return_value.json.return_value = orders
mock_client.get_orders_for_all_linked_accounts.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), 0)

Expand All @@ -56,7 +60,7 @@ def test_success_no_account_id(
@patch('schwab.scripts.orders_codegen.client_from_token_file')
@patch('schwab.scripts.orders_codegen.construct_repeat_order')
@patch('schwab.scripts.orders_codegen.code_for_builder')
def test_no_account_id_no_such_order(
def test_no_account_id_no_recent_orders(
self,
mock_code_for_builder,
mock_construct_repeat_order,
Expand All @@ -66,12 +70,15 @@ def test_no_account_id_no_such_order(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')

orders = []

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_query.return_value.json.return_value = orders
mock_client.get_orders_for_all_linked_accounts.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), 0)

Expand All @@ -94,18 +101,21 @@ def test_no_account_error(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')

orders = {'error': 'invalid'}

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_query.return_value.json.return_value = orders
mock_client.get_orders_for_all_linked_accounts.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), -1)

mock_construct_repeat_order.assert_not_called()
mock_print.assert_called_once_with(
AnyStringWith('TDA returned error: "invalid"'))
AnyStringWith('Schwab returned error: "invalid"'))


@no_duplicates
Expand All @@ -123,6 +133,8 @@ def test_success_account_id(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')
self.add_arg('--account_id')
self.add_arg('123456')

Expand All @@ -135,10 +147,18 @@ def test_success_account_id(

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_path.return_value.json.return_value = orders
mock_client.get_account_numbers.return_value = httpx.Response(
200,
json=[{
'accountNumber': '123456',
'hashValue': 'hash-value',
}])
mock_client.get_orders_for_account.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), 0)

mock_client.get_account_numbers.assert_called_once()
mock_construct_repeat_order.assert_called_once_with(orders[3])
mock_print.assert_has_calls([
call('# Order ID', 401),
Expand All @@ -160,14 +180,23 @@ def test_account_id_no_orders(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')
self.add_arg('--account_id')
self.add_arg('123456')

orders = []

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_path.return_value.json.return_value = orders
mock_client.get_account_numbers.return_value = httpx.Response(
200,
json=[{
'accountNumber': '123456',
'hashValue': 'hash-value',
}])
mock_client.get_orders_for_account.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), 0)

Expand All @@ -190,21 +219,67 @@ def test_account_id_error(
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')
self.add_arg('--account_id')
self.add_arg('123456')

orders = {'error': 'invalid'}

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_by_path.return_value.json.return_value = orders
mock_client.get_account_numbers.return_value = httpx.Response(
200,
json=[{
'accountNumber': '123456',
'hashValue': 'hash-value',
}])
mock_client.get_orders_for_account.return_value \
= httpx.Response(200, json={'error': 'invalid'})

self.assertEqual(self.main(), -1)

mock_construct_repeat_order.assert_not_called
mock_print.assert_called_once_with(
AnyStringWith('TDA returned error: "invalid"'))
AnyStringWith('Schwab returned error: "invalid"'))


@no_duplicates
@patch('builtins.print')
@patch('schwab.scripts.orders_codegen.client_from_token_file')
@patch('schwab.scripts.orders_codegen.construct_repeat_order')
@patch('schwab.scripts.orders_codegen.code_for_builder')
def test_success_account_hash(
self,
mock_code_for_builder,
mock_construct_repeat_order,
mock_client_from_token_file,
mock_print):
self.add_arg('--token_file')
self.add_arg('filename.json')
self.add_arg('--api_key')
self.add_arg('api-key')
self.add_arg('--app_secret')
self.add_arg('app-secret')
self.add_arg('--account_hash')
self.add_arg('account-hash')

orders = [
{'orderId': 201},
{'orderId': 101},
{'orderId': 301},
{'orderId': 401},
]

mock_client = MagicMock()
mock_client_from_token_file.return_value = mock_client
mock_client.get_orders_for_account.return_value \
= httpx.Response(200, json=orders)

self.assertEqual(self.main(), 0)

mock_construct_repeat_order.assert_called_once_with(orders[3])
mock_print.assert_has_calls([
call('# Order ID', 401),
call(mock_code_for_builder.return_value)])

class ScriptInvocationTest(unittest.TestCase):

Expand Down

0 comments on commit fce52ee

Please sign in to comment.