fixed tests and cleaned up

ani · ani · commit 7036e7031092 · 2024-02-20T14:57:40.000-05:00
diff --git a/examples/read_hermes.py b/examples/read_hermes.py
@@ -8,28 +8,29 @@ async def get_hermes_prices():
     hermes_client = HermesClient([])
     feed_ids = await hermes_client.get_price_feed_ids()
     feed_ids_rel = feed_ids[:2]
-    version = 1
+    version_http = 2
+    version_ws = 1
 
     hermes_client.add_feed_ids(feed_ids_rel)
 
-    prices_latest = await hermes_client.get_all_prices(version=version)
+    prices_latest = await hermes_client.get_all_prices(version=version_http)
 
     sd = list(prices_latest.keys())[0]
 
     for feed_id, price_feed in prices_latest.items():
         print("Initial prices")
-        price_latest = price_feed["price"].price
-        conf_latest = price_feed["price"].conf
-        print(f"Feed ID: {feed_id}, Price: {price_latest}, Confidence: {conf_latest}, Time: {price_feed['price'].publish_time}")
+        print(f"Feed ID: {feed_id}, Price: {price_feed['price'].price}, Confidence: {price_feed['price'].conf}, Time: {price_feed['price'].publish_time}")
 
     print("Starting web socket...")
-    ws_call = hermes_client.ws_pyth_prices(version=version)
-    asyncio.create_task(ws_call)
+    ws_call = hermes_client.ws_pyth_prices(version=version_ws)
+    ws_task = asyncio.create_task(ws_call)
 
     while True:
         await asyncio.sleep(5)
+        if ws_task.done():
+            break
         print("Latest prices:")
         for feed_id, price_feed in hermes_client.prices_dict.items():
-            print(f"Feed ID: {feed_id}, Price: {price_latest}, Confidence: {conf_latest}, Time: {price_feed['price'].publish_time}")
+            print(f"Feed ID: {feed_id}, Price: {price_feed['price'].price}, Confidence: {price_feed['price'].conf}, Time: {price_feed['price'].publish_time}")
 
 asyncio.run(get_hermes_prices())
diff --git a/pythclient/hermes.py b/pythclient/hermes.py
@@ -33,7 +33,7 @@ def __init__(self, feed_ids: list[str], endpoint=HERMES_ENDPOINT_HTTPS, ws_endpo
         self.client = httpx.AsyncClient()
         self.endpoint = endpoint
         self.ws_endpoint = ws_endpoint
-        self.feed_batch_size = feed_batch_size
+        self.feed_batch_size = feed_batch_size # max number of feed IDs to query at once in https requests
 
     async def get_price_feed_ids(self) -> list[str]:
         """
@@ -70,7 +70,7 @@ def extract_price_feed_v1(data: dict) -> PriceFeed:
     @staticmethod
     def extract_price_feed_v2(data: dict) -> list[PriceFeed]:
         """
-        Extracts PriceFeed objects from the v2 JSON response (multiple price feeds) from Hermes.
+        Extracts PriceFeed objects from the v2 JSON response (array of price feeds) from Hermes.
         """
         update_data = data["binary"]["data"]
 
@@ -140,7 +140,7 @@ async def get_all_prices(self, version=2) -> dict[str, PriceFeed]:
         """
         Queries the Hermes http endpoint for the latest price feeds for all feed IDs in the class object.
 
-        There are limitations on the number of feed IDs that can be queried at once, so this function queries the feed IDs in batches.
+        There is a limit on the number of feed IDs that can be queried at once, so this function queries the feed IDs in batches.
         """
         pyth_prices_latest = []
         i = 0
@@ -159,7 +159,7 @@ async def ws_pyth_prices(self, version=1):
         """
         if version != 1:
             parse_unsupported_version(version)
-
+        
         async with websockets.connect(self.ws_endpoint) as ws:
             while True:
                 # add new price feed ids to the ws subscription
@@ -187,34 +187,4 @@ async def ws_pyth_prices(self, version=1):
                     self.prices_dict[feed_id] = self.extract_price_feed_v1(new_feed)
 
                 except Exception as e:
-                    raise Exception(f"Error in price_update message: {msg}") from e
-
-
-async def main():
-    hermes_client = HermesClient([])
-    feed_ids = await hermes_client.get_price_feed_ids()
-    feed_ids_rel = feed_ids[:50]
-    version = 2
-
-    hermes_client.add_feed_ids(feed_ids_rel)
-    
-    prices_latest = await hermes_client.get_pyth_prices_latest(feed_ids[:50], version=version)
-
-    try:
-        price_at_time = await hermes_client.get_pyth_price_at_time(feed_ids[0], 1_700_000_000, version=version)
-        print(price_at_time)
-    except Exception as e:
-        print(f"Error in get_pyth_price_at_time, {e}")
-
-    all_prices = await hermes_client.get_all_prices(version=version)
-
-    print("Starting web socket...")
-    ws_call = hermes_client.ws_pyth_prices(version=version)
-    asyncio.create_task(ws_call)
-
-    while True:
-        await asyncio.sleep(1)
-
-
-if __name__ == "__main__":
-    asyncio.run(main())
+                    raise Exception(f"Error in price_update message: {msg}") from e
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name='pythclient',
-    version='0.1.22',
+    version='0.1.23',
     packages=['pythclient'],
     author='Pyth Developers',
     author_email='contact@pyth.network',
diff --git a/tests/test_hermes.py b/tests/test_hermes.py
@@ -50,24 +50,24 @@ def data_v2():
         },
         "parsed": [
             {
-            "id": "e62df6c8b4a85fe1a67db44dc12de5db330f7ac66b72dc658afedf0f4a415b43",
-            "price": {
-                "price": "5190075917635",
-                "conf": "2661582364",
-                "expo": -8,
-                "publish_time": 1708363256
-            },
-            "ema_price": {
-                "price": "5209141800000",
-                "conf": "3290086600",
-                "expo": -8,
-                "publish_time": 1708363256
-            },
-            "metadata": {
-                "slot": 125976528,
-                "proof_available_time": 1708363257,
-                "prev_publish_time": 1708363255
-            }
+                "id": "e62df6c8b4a85fe1a67db44dc12de5db330f7ac66b72dc658afedf0f4a415b43",
+                "price": {
+                    "price": "5190075917635",
+                    "conf": "2661582364",
+                    "expo": -8,
+                    "publish_time": 1708363256
+                },
+                "ema_price": {
+                    "price": "5209141800000",
+                    "conf": "3290086600",
+                    "expo": -8,
+                    "publish_time": 1708363256
+                },
+                "metadata": {
+                    "slot": 125976528,
+                    "proof_available_time": 1708363257,
+                    "prev_publish_time": 1708363255
+                }
             }
         ]
     }
@@ -79,7 +79,7 @@ def mock_get_price_feed_ids(mocker: MockerFixture):
     return async_mock
 
 @pytest.mark.asyncio
-async def test_hermes_add_feed_ids(hermes_client: HermesClient):
+async def test_hermes_add_feed_ids(hermes_client: HermesClient, mock_get_price_feed_ids: AsyncMock):
     mock_get_price_feed_ids.return_value = ["ff61491a931112ddf1bd8147cd1b641375f79f5825126d665480874634fd0ace"]
 
     feed_ids = await hermes_client.get_price_feed_ids()
@@ -89,10 +89,8 @@ async def test_hermes_add_feed_ids(hermes_client: HermesClient):
 
     hermes_client.add_feed_ids(feed_ids)
 
-    assert hermes_client.feed_ids == list(set(feed_ids_pre + feed_ids))
-    assert hermes_client.pending_feed_ids == list(set(pending_feed_ids_pre + feed_ids))
-
-
+    assert set(hermes_client.feed_ids) == set(feed_ids_pre + feed_ids)
+    assert set(hermes_client.pending_feed_ids) == set(pending_feed_ids_pre + feed_ids)
 
 def test_hermes_extract_price_feed_v1(hermes_client: HermesClient, data_v1: dict):
     price_feed = hermes_client.extract_price_feed_v1(data_v1)
@@ -103,5 +101,6 @@ def test_hermes_extract_price_feed_v1(hermes_client: HermesClient, data_v1: dict
 def test_hermes_extract_price_feed_v2(hermes_client: HermesClient, data_v2: dict):
     price_feed = hermes_client.extract_price_feed_v2(data_v2)
 
-    assert isinstance(price_feed, dict)
-    assert set(price_feed.keys()) == set(PriceFeed.__annotations__.keys())
+    assert isinstance(price_feed, list)
+    assert isinstance(price_feed[0], dict)
+    assert set(price_feed[0].keys()) == set(PriceFeed.__annotations__.keys())
