Format

drift-labs · Dec 16, 2024 · 7be1374 · 7be1374
1 parent 05c4560
commit 7be1374
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/examples/fetch_all_markets.py b/examples/fetch_all_markets.py
@@ -1,16 +1,14 @@
 import asyncio
 import os
 
-from anchorpy import Provider
-from anchorpy import Wallet
-from driftpy.drift_client import AccountSubscriptionConfig
-from driftpy.drift_client import DriftClient
+from anchorpy.provider import Provider, Wallet
 from solana.rpc.async_api import AsyncClient
 from solders.keypair import Keypair
 
+from driftpy.drift_client import AccountSubscriptionConfig, DriftClient
+
 
 async def get_all_market_names():
-    env = "mainnet-beta"  # 'devnet'
     rpc = os.environ.get("MAINNET_RPC_ENDPOINT")
     kp = Keypair()  # random wallet
     wallet = Wallet(kp)
@@ -19,10 +17,10 @@ async def get_all_market_names():
     drift_client = DriftClient(
         provider.connection,
         provider.wallet,
-        env.split("-")[0],
-        account_subscription=AccountSubscriptionConfig("cached"),
+        "mainnet",
+        account_subscription=AccountSubscriptionConfig("websocket"),
     )
-
+    await drift_client.subscribe()
     all_perps_markets = await drift_client.program.account["PerpMarket"].all()
     sorted_all_perps_markets = sorted(
         all_perps_markets, key=lambda x: x.account.market_index
@@ -46,10 +44,26 @@ async def get_all_market_names():
         print(market)
 
     result = result_perp + result_spot[1:]
+
+    print("Here are some prices:")
+    print(drift_client.get_oracle_price_data_for_perp_market(0))
+    print(drift_client.get_oracle_price_data_for_spot_market(0))
+    await drift_client.unsubscribe()
     return result
 
 
 if __name__ == "__main__":
-    loop = asyncio.new_event_loop()
-    answer = loop.run_until_complete(get_all_market_names())
-    print(answer)
+    try:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        answer = loop.run_until_complete(get_all_market_names())
+        print(answer)
+    finally:
+        # Clean up pending tasks
+        pending = asyncio.all_tasks(loop)
+        for task in pending:
+            task.cancel()
+
+        # Run loop until tasks complete/cancel
+        loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+        loop.close()
diff --git a/tests/dlob_test_constants.py b/tests/dlob_test_constants.py
@@ -1,11 +1,13 @@
+from solders.pubkey import Pubkey
+
+from driftpy.constants.config import devnet_spot_market_configs
 from driftpy.constants.numeric_constants import (
     AMM_TO_QUOTE_PRECISION_RATIO,
     BASE_PRECISION,
     PRICE_PRECISION,
     QUOTE_PRECISION,
     SPOT_CUMULATIVE_INTEREST_PRECISION,
     SPOT_MARKET_CUMULATIVE_INTEREST_PRECISION,
-    SPOT_MARKET_CUMULATIVE_INTEREST_PRECISION_EXP,
     SPOT_MARKET_WEIGHT_PRECISION,
 )
 from driftpy.types import (
@@ -23,8 +25,6 @@
     PoolBalance,
     SpotMarketAccount,
 )
-from driftpy.constants.config import devnet_spot_market_configs
-from solders.pubkey import Pubkey
 
 mock_pool_balance = PoolBalance(
     scaled_balance=0, market_index=0, padding=[0] * 6