From 37d6447aa799daea57be4e4f39895383211a0abf Mon Sep 17 00:00:00 2001 From: rcholic Date: Fri, 7 Jun 2024 18:20:19 -0700 Subject: [PATCH] unit test get accounts --- cschwabpy/SchwabAsyncClient.py | 53 +++++++++++- cschwabpy/models/trade_models.py | 8 ++ tests/data/mock_schwab_api_resp.json | 119 +++++++++++++++++++++++++++ tests/test_models.py | 68 ++++++++++++++- 4 files changed, 245 insertions(+), 3 deletions(-) diff --git a/cschwabpy/SchwabAsyncClient.py b/cschwabpy/SchwabAsyncClient.py index 19c89ee..5879438 100644 --- a/cschwabpy/SchwabAsyncClient.py +++ b/cschwabpy/SchwabAsyncClient.py @@ -6,7 +6,7 @@ OptionExpiration, OptionExpirationChainResponse, ) -from cschwabpy.models.trade_models import AccountNumberModel +from cschwabpy.models.trade_models import AccountNumberModel, SecuritiesAccount, Account from typing import Optional, List, Mapping from cschwabpy.costants import ( SCHWAB_API_BASE_URL, @@ -113,7 +113,6 @@ async def get_account_numbers_async(self) -> List[AccountNumberModel]: url=target_url, params={}, headers=self.__auth_header() ) json_res = response.json() - print("json_res: ", json_res) account_numbers: List[AccountNumberModel] = [] for account_json in json_res: account_numbers.append(AccountNumberModel(**account_json)) @@ -122,6 +121,56 @@ async def get_account_numbers_async(self) -> List[AccountNumberModel]: if not self.__keep_client_alive: await client.aclose() + async def get_accounts_async( + self, includ_positions: bool = True, with_account_number: Optional[str] = None + ) -> List[Account]: + """get all accounts except a specific account_number is provided.""" + await self._ensure_valid_access_token() + target_url = f"{SCHWAB_TRADER_API_BASE_URL}/accounts" + if with_account_number is not None: + target_url = f"{target_url}/{with_account_number}" + + if includ_positions: + target_url = f"{target_url}?fields=positions" + + client = httpx.AsyncClient() if self.__client is None else self.__client + try: + response = await client.get( + url=target_url, params={}, headers=self.__auth_header() + ) + if response.status_code == 200: + json_res = response.json() + if with_account_number is None: + accounts: List[SecuritiesAccount] = [] + for account_json in json_res: + securities_account = SecuritiesAccount( + **account_json + ).securitiesAccount + accounts.append(securities_account) + return accounts + else: + securities_account = SecuritiesAccount(**json_res).securitiesAccount + return [securities_account] + else: + raise Exception( + "Failed to get accounts. Status: ", response.status_code + ) + finally: + if not self.__keep_client_alive: + await client.aclose() + + async def get_single_account_async( + self, with_account_number: str, include_positions: bool = True + ) -> Optional[Account]: + """Convenience method to get a single account by account number.""" + account = await self.get_accounts_async( + includ_positions=include_positions, with_account_number=with_account_number + ) + if account is None or len(account) == 0: + return None + + return account[0] + async def get_option_expirations_async( self, underlying_symbol: str ) -> List[OptionExpiration]: diff --git a/cschwabpy/models/trade_models.py b/cschwabpy/models/trade_models.py index dd878a5..3fc65f5 100644 --- a/cschwabpy/models/trade_models.py +++ b/cschwabpy/models/trade_models.py @@ -95,6 +95,14 @@ class Account(JSONSerializableBaseModel): currentBalances: Optional[MarginBalance] = None projectedBalances: Optional[MarginBalance] = None + @property + def is_margin(self) -> bool: + return self.type_ == AccountType.MARGIN + + @property + def is_cash(self) -> bool: + return self.type_ == AccountType.CASH + class MarginAccount(Account): type_: AccountType = Field(AccountType.MARGIN, alias="type") diff --git a/tests/data/mock_schwab_api_resp.json b/tests/data/mock_schwab_api_resp.json index 6a7d1bf..828df5d 100644 --- a/tests/data/mock_schwab_api_resp.json +++ b/tests/data/mock_schwab_api_resp.json @@ -1,4 +1,122 @@ { + "single_account": { + "securitiesAccount": { + "accountNumber": "123", + "roundTrips": 0, + "isDayTrader": false, + "isClosingOnlyRestricted": false, + "pfcbFlag": false, + "type": "MARGIN", + "positions": [ + { + "shortQuantity": 0, + "averagePrice": 0, + "currentDayProfitLoss": 0, + "currentDayProfitLossPercentage": 0, + "longQuantity": 0, + "settledLongQuantity": 0, + "settledShortQuantity": 0, + "agedQuantity": 0, + "instrument": { + "cusip": "123", + "symbol": "spy", + "description": "sfsdfsadf", + "instrumentId": 0, + "netChange": 0, + "type": "SWEEP_VEHICLE" + }, + "marketValue": 0, + "maintenanceRequirement": 0, + "averageLongPrice": 0, + "averageShortPrice": 0, + "taxLotAverageLongPrice": 0, + "taxLotAverageShortPrice": 0, + "longOpenProfitLoss": 0, + "shortOpenProfitLoss": 0, + "previousSessionLongQuantity": 0, + "previousSessionShortQuantity": 0, + "currentDayCost": 0 + } + ], + "initialBalances": { + "accruedInterest": 0, + "availableFundsNonMarginableTrade": 0, + "bondValue": 0, + "buyingPower": 0, + "cashBalance": 0, + "cashAvailableForTrading": 0, + "cashReceipts": 0, + "dayTradingBuyingPower": 0, + "dayTradingBuyingPowerCall": 0, + "dayTradingEquityCall": 0, + "equity": 0, + "equityPercentage": 0, + "liquidationValue": 0, + "longMarginValue": 0, + "longOptionMarketValue": 0, + "longStockValue": 0, + "maintenanceCall": 0, + "maintenanceRequirement": 0, + "margin": 0, + "marginEquity": 0, + "moneyMarketFund": 0, + "mutualFundValue": 0, + "regTCall": 0, + "shortMarginValue": 0, + "shortOptionMarketValue": 0, + "shortStockValue": 0, + "totalCash": 0, + "isInCall": 0, + "unsettledCash": 0, + "pendingDeposits": 0, + "marginBalance": 0, + "shortBalance": 0, + "accountValue": 0 + }, + "currentBalances": { + "availableFunds": 0, + "availableFundsNonMarginableTrade": 0, + "buyingPower": 0, + "buyingPowerNonMarginableTrade": 0, + "dayTradingBuyingPower": 0, + "dayTradingBuyingPowerCall": 0, + "equity": 0, + "equityPercentage": 0, + "longMarginValue": 0, + "maintenanceCall": 0, + "maintenanceRequirement": 0, + "marginBalance": 0, + "regTCall": 0, + "shortBalance": 0, + "shortMarginValue": 0, + "sma": 0, + "isInCall": 0, + "stockBuyingPower": 0, + "optionBuyingPower": 0 + }, + "projectedBalances": { + "availableFunds": 0, + "availableFundsNonMarginableTrade": 0, + "buyingPower": 0, + "buyingPowerNonMarginableTrade": 0, + "dayTradingBuyingPower": 0, + "dayTradingBuyingPowerCall": 0, + "equity": 0, + "equityPercentage": 0, + "longMarginValue": 0, + "maintenanceCall": 0, + "maintenanceRequirement": 0, + "marginBalance": 0, + "regTCall": 0, + "shortBalance": 0, + "shortMarginValue": 0, + "sma": 0, + "isInCall": 0, + "stockBuyingPower": 0, + "optionBuyingPower": 0 + } + } + }, "securities_account": [ { "securitiesAccount": { @@ -7,6 +125,7 @@ "isDayTrader": false, "isClosingOnlyRestricted": false, "pfcbFlag": false, + "type": "MARGIN", "positions": [ { "shortQuantity": 0, diff --git a/tests/test_models.py b/tests/test_models.py index 6533e12..a360841 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -11,7 +11,12 @@ OptionContractType, OptionContractStrategy, ) -from cschwabpy.models.trade_models import SecuritiesAccount, MarginAccount, CashAccount +from cschwabpy.models.trade_models import ( + SecuritiesAccount, + MarginAccount, + CashAccount, + AccountType, +) from cschwabpy.models.token import Tokens, LocalTokenStore from cschwabpy.SchwabAsyncClient import SchwabAsyncClient @@ -69,6 +74,67 @@ def test_parsing_securities_account(): assert len(accounts) == 1 +@pytest.mark.asyncio +async def test_get_single_account(httpx_mock: HTTPXMock): + json_mock = get_mock_response()["single_account"] + mocked_token = mock_tokens() + token_store = LocalTokenStore() + if os.path.exists(Path(token_store.token_output_path)): + os.remove(token_store.token_output_path) # clean up before test + + token_store.save_tokens(mocked_token) + symbol = "$SPX" + httpx_mock.add_response(json=json_mock) + async with httpx.AsyncClient() as client: + cschwab_client = SchwabAsyncClient( + app_client_id="fake_id", + app_secret="fake_secret", + token_store=token_store, + tokens=mocked_token, + http_client=client, + ) + single_account = await cschwab_client.get_single_account_async( + with_account_number="123", include_positions=True + ) + assert single_account is not None + assert single_account.accountNumber == "123" + assert single_account.type_ == AccountType.MARGIN + assert len(single_account.positions) == 1 + assert single_account.is_margin + + +@pytest.mark.asyncio +async def test_get_securities_account(httpx_mock: HTTPXMock): + json_mock = get_mock_response()["securities_account"] # single_account + mocked_token = mock_tokens() + token_store = LocalTokenStore() + if os.path.exists(Path(token_store.token_output_path)): + os.remove(token_store.token_output_path) # clean up before test + + token_store.save_tokens(mocked_token) + symbol = "$SPX" + httpx_mock.add_response(json=json_mock) + async with httpx.AsyncClient() as client: + cschwab_client = SchwabAsyncClient( + app_client_id="fake_id", + app_secret="fake_secret", + token_store=token_store, + tokens=mocked_token, + http_client=client, + ) + securities_accounts = await cschwab_client.get_accounts_async( + includ_positions=True + ) + + assert securities_accounts is not None + assert len(securities_accounts) == 1 + assert securities_accounts[0].accountNumber == "123" + assert securities_accounts[0].type_ == AccountType.MARGIN + assert securities_accounts[0].type_ == AccountType.MARGIN + assert len(securities_accounts[0].positions) == 1 + assert securities_accounts[0].is_margin + + @pytest.mark.asyncio async def test_download_option_chain(httpx_mock: HTTPXMock): mock_option_chain_resp = get_mock_response()