frank/update DriftUser oracle price calls (#131)

drift-labs · Mar 5, 2024 · d167ce1 · d167ce1
1 parent 5744a67
commit d167ce1
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 10 deletions.
diff --git a/src/driftpy/accounts/oracle.py b/src/driftpy/accounts/oracle.py
@@ -119,8 +119,11 @@ def decode_swb_price_info(data: bytes):
 
     price = convert_switchboard_decimal(round.result.mantissa, round.result.scale)
 
-    conf = convert_switchboard_decimal(
-        round.std_deviation.mantissa, round.std_deviation.scale
+    conf = max(
+        convert_switchboard_decimal(
+            round.std_deviation.mantissa, round.std_deviation.scale
+        ),
+        (price // 1_000),
     )
 
     has_sufficient_number_of_data_points = (
@@ -130,7 +133,7 @@ def decode_swb_price_info(data: bytes):
     slot = round.round_open_slot
 
     return OraclePriceData(
-        price, slot, conf, 1, 1, has_sufficient_number_of_data_points
+        price, slot, conf, None, None, has_sufficient_number_of_data_points
     )
 
 

diff --git a/src/driftpy/drift_user.py b/src/driftpy/drift_user.py
@@ -153,7 +153,7 @@ def get_perp_market_liability(
 
             market = self.drift_client.get_perp_market_account(position.market_index)
 
-            price = self.drift_client.get_oracle_price_data(market.amm.oracle).price
+            price = (self.get_oracle_data_for_perp_market(position.market_index)).price
             base_asset_amount = (
                 calculate_worst_case_base_asset_amount(position)
                 if include_open_orders
@@ -288,8 +288,8 @@ def calculate_weighted_perp_position_value(
                 market.quote_spot_market_index
             )
 
-            quote_oracle_price_data = self.drift_client.get_oracle_price_data(
-                quote_spot_market.oracle
+            quote_oracle_price_data = self.get_oracle_data_for_spot_market(
+                QUOTE_SPOT_MARKET_INDEX
             )
 
             if strict:
@@ -536,8 +536,8 @@ def get_spot_market_asset_and_liability_value(
             spot_market_account = self.drift_client.get_spot_market_account(
                 spot_position.market_index
             )
-            oracle_price_data = self.drift_client.get_oracle_price_data(
-                spot_market_account.oracle
+            oracle_price_data = self.get_oracle_data_for_spot_market(
+                spot_position.market_index
             )
 
             twap_5m = None
@@ -963,7 +963,7 @@ def calculate_free_collateral_delta_for_spot(
         if signed_token_amount > 0:
             asset_weight = calculate_asset_weight(
                 signed_token_amount,
-                self.drift_client.get_oracle_price_data(market.oracle).price,
+                self.get_oracle_data_for_spot_market(market.market_index).price,
                 market,
                 MarginCategory.MAINTENANCE,
             )
@@ -1091,7 +1091,7 @@ def get_spot_liq_price(
             case _:
                 raise Exception(f"Invalid balance type: {position.balance_type}")
 
-        price = self.drift_client.get_oracle_price_data(spot_market.oracle).price
+        price = self.get_oracle_data_for_spot_market(spot_market.market_index).price
         liq_price = price + liq_price_delta
         liq_price /= PRICE_PRECISION
 

diff --git a/tests/math/helpers.py b/tests/math/helpers.py
@@ -168,9 +168,23 @@ def get_oracle(key):
             has_sufficient_number_of_data_points=True,
         )
 
+    def get_oracle_price_data_for_perp_market(market_index):
+        market = get_perp(market_index)
+        return get_oracle(market.amm.oracle)
+
+    def get_oracle_price_data_for_spot_market(market_index):
+        market = get_spot(market_index)
+        return get_oracle(market.oracle)
+
     muser.get_user_account = get_user
     muser.drift_client.get_perp_market_account = get_perp
     muser.drift_client.get_spot_market_account = get_spot
     muser.drift_client.get_oracle_price_data = get_oracle
+    muser.drift_client.get_oracle_price_data_for_perp_market = (
+        get_oracle_price_data_for_perp_market
+    )
+    muser.drift_client.get_oracle_price_data_for_spot_market = (
+        get_oracle_price_data_for_spot_market
+    )
 
     return muser