From beb6280e161a3ab3b3ee3c7d2adc8471c764e397 Mon Sep 17 00:00:00 2001
From: rcholic <ivytony@gmail.com>
Date: Fri, 7 Jun 2024 19:35:00 -0700
Subject: [PATCH] test api call for order retrieval

---
 cschwabpy/SchwabAsyncClient.py | 43 +++++++++++++++++++++++++++++++++-
 cschwabpy/util.py              |  6 +++++
 tests/test_models.py           | 34 +++++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/cschwabpy/SchwabAsyncClient.py b/cschwabpy/SchwabAsyncClient.py
index a5eca7f..6760203 100644
--- a/cschwabpy/SchwabAsyncClient.py
+++ b/cschwabpy/SchwabAsyncClient.py
@@ -6,7 +6,16 @@
     OptionExpiration,
     OptionExpirationChainResponse,
 )
-from cschwabpy.models.trade_models import AccountNumberModel, SecuritiesAccount, Account
+from cschwabpy.models.trade_models import (
+    AccountNumberModel,
+    SecuritiesAccount,
+    Account,
+    OrderStatus,
+    Order,
+)
+import cschwabpy.util as util
+
+from datetime import datetime, timedelta
 from typing import Optional, List, Mapping
 from cschwabpy.costants import (
     SCHWAB_API_BASE_URL,
@@ -171,6 +180,38 @@ async def get_single_account_async(
 
         return account[0]
 
+    async def get_orders_async(
+        self,
+        account_number: str,
+        from_entered_time: datetime,
+        to_entered_time: datetime,
+        max_count: int = 1000,
+        status: Optional[OrderStatus] = None,
+    ) -> List[Order]:
+        await self._ensure_valid_access_token()
+        target_url = f"{SCHWAB_TRADER_API_BASE_URL}/accounts/{account_number}/orders"
+        target_url += f"?fromEnteredTime={util.to_iso8601_str(from_entered_time)}&toEnteredTime={util.to_iso8601_str(to_entered_time)}&maxResults={max_count}"
+        if status is not None:
+            target_url += f"&status={status.value}"
+
+        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()
+                orders: List[Order] = []
+                for order_json in json_res:
+                    order = Order(**order_json)
+                    orders.append(order)
+                return orders
+            else:
+                raise Exception("Failed to get orders. Status: ", response.status_code)
+        finally:
+            if not self.__keep_client_alive:
+                await client.aclose()
+
     async def get_option_expirations_async(
         self, underlying_symbol: str
     ) -> List[OptionExpiration]:
diff --git a/cschwabpy/util.py b/cschwabpy/util.py
index b51ac38..f1f32fc 100644
--- a/cschwabpy/util.py
+++ b/cschwabpy/util.py
@@ -36,6 +36,12 @@ def ts_to_date_string(
     return ts_to_datetime(ts, tz).strftime(YMD_FMT)
 
 
+def to_iso8601_str(dt: datetime) -> str:
+    dt = dt.astimezone(eastern_tz)
+    dt_str = dt.strftime("%Y-%m-%dT%H:%M:%S")
+    return dt_str + ".000Z"
+
+
 def today_str(tz: pytz.BaseTzInfo = eastern_tz) -> str:  # type: ignore
     """Today in string. Returns Y-m-d."""  # noqa: DAR201
     return now(tz=tz).strftime(YMD_FMT)
diff --git a/tests/test_models.py b/tests/test_models.py
index 6c108e9..25afeab 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -3,6 +3,7 @@
 import typing
 import httpx
 import pytest
+from datetime import datetime, timedelta
 from pytest_httpx import HTTPXMock
 from pathlib import Path
 from cschwabpy.models import (
@@ -17,6 +18,7 @@
     CashAccount,
     AccountType,
     Order,
+    OrderStatus,
 )
 from cschwabpy.models.token import Tokens, LocalTokenStore
 from cschwabpy.SchwabAsyncClient import SchwabAsyncClient
@@ -83,6 +85,38 @@ def test_parsing_order():
     assert order_obj.cancelable == False
 
 
+@pytest.mark.asyncio
+async def test_get_order(httpx_mock: HTTPXMock):
+    json_mock = get_mock_response()["single_order"]
+    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)
+    httpx_mock.add_response(json=[json_mock])  # make it array type
+    from_entered_time = datetime.now() - timedelta(hours=3)
+    to_entered_time = datetime.now()
+    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,
+        )
+        retrieved_orders = await cschwab_client.get_orders_async(
+            account_number="123",
+            from_entered_time=from_entered_time,
+            to_entered_time=to_entered_time,
+            status=OrderStatus.FILLED,
+        )
+        assert retrieved_orders is not None
+        assert len(retrieved_orders) == 1
+        assert retrieved_orders[0].orderId == 456
+        assert retrieved_orders[0].cancelable == False
+
+
 @pytest.mark.asyncio
 async def test_get_single_account(httpx_mock: HTTPXMock):
     json_mock = get_mock_response()["single_account"]