Consolidate spot_price_after_open_short fns, add test

crates/hyperdrive-math/src/short/open.rs

@@ -52,22 +52,21 @@ impl State {
     /// Calculates the spot price after opening a Hyperdrive short.
     pub fn calculate_spot_price_after_short(
         &self,
-        short_amount: FixedPoint,
-        spot_price: FixedPoint,
-        open_vault_share_price: FixedPoint,
-    ) -> Result<FixedPoint> {
-        self.calculate_open_short(short_amount, spot_price, open_vault_share_price)
-            .map(|short_deposit| self.spot_price_after_short(short_deposit, short_amount))
-    }
-
-    fn spot_price_after_short(
-        &self,
-        base_amount: FixedPoint,
         bond_amount: FixedPoint,
+        base_amount: Option<FixedPoint>,
     ) -> FixedPoint {
+        let shares_amount = match base_amount {
+            Some(base_amount) => base_amount / self.vault_share_price(),
+            None => {
+                let spot_price = self.calculate_spot_price();
+                self.calculate_shares_out_given_bonds_in_down(bond_amount)
+                    - self.open_short_curve_fee(bond_amount, spot_price)
+                    + self.open_short_governance_fee(bond_amount, spot_price)
+            }
+        };
         let mut state: State = self.clone();
         state.info.bond_reserves += bond_amount.into();
-        state.info.share_reserves -= (base_amount / state.vault_share_price()).into();
+        state.info.share_reserves -= shares_amount.into();
         state.calculate_spot_price()
     }
 
@@ -86,3 +85,83 @@ impl State {
         self.calculate_shares_out_given_bonds_in_down_safe(short_amount)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use eyre::Result;
+    use fixed_point_macros::fixed;
+    use rand::{thread_rng, Rng};
+    use test_utils::{
+        agent::Agent,
+        chain::{Chain, TestChain},
+        constants::FUZZ_RUNS,
+    };
+
+    use super::*;
+
+    #[tokio::test]
+    async fn fuzz_calculate_spot_price_after_short() -> Result<()> {
+        // Spawn a test chain and create two agents -- Alice and Bob. Alice is
+        // funded with a large amount of capital so that she can initialize the
+        // pool. Bob is funded with a small amount of capital so that we can
+        // test opening a short and verify that the ending spot price is what we
+        // expect.
+        let mut rng = thread_rng();
+        let chain = TestChain::new(2).await?;
+        let (alice, bob) = (chain.accounts()[0].clone(), chain.accounts()[1].clone());
+        let mut alice =
+            Agent::new(chain.client(alice).await?, chain.addresses().clone(), None).await?;
+        let mut bob = Agent::new(chain.client(bob).await?, chain.addresses(), None).await?;
+
+        for _ in 0..*FUZZ_RUNS {
+            // Snapshot the chain.
+            let id = chain.snapshot().await?;
+
+            // Fund Alice and Bob.
+            let fixed_rate = rng.gen_range(fixed!(0.01e18)..=fixed!(0.1e18));
+            let contribution = rng.gen_range(fixed!(10_000e18)..=fixed!(500_000_000e18));
+            let budget = rng.gen_range(fixed!(10e18)..=fixed!(500_000_000e18));
+            alice.fund(contribution).await?;
+            bob.fund(budget).await?;
+
+            // Alice initializes the pool.
+            alice.initialize(fixed_rate, contribution, None).await?;
+
+            // Attempt to predict the spot price after opening a short.
+            let short_amount =
+                rng.gen_range(fixed!(0.01e18)..=bob.calculate_max_short(None).await?);
+            let current_state = bob.get_state().await?;
+            let expected_spot_price =
+                current_state.calculate_spot_price_after_short(short_amount, None);
+
+            // Open the short.
+            bob.open_short(short_amount, None, None).await?;
+
+            // Verify that the predicted spot price is equal to the ending spot
+            // price. These won't be exactly equal because the vault share price
+            // increases between the prediction and opening the short.
+            let actual_spot_price = bob.get_state().await?.calculate_spot_price();
+            let delta = if actual_spot_price > expected_spot_price {
+                actual_spot_price - expected_spot_price
+            } else {
+                expected_spot_price - actual_spot_price
+            };
+            // TODO: Why can't this pass with a tolerance of 1e9?
+            let tolerance = fixed!(1e10);
+
+            assert!(
+                delta < tolerance,
+                "expected: delta = {} < {} = tolerance",
+                delta,
+                tolerance
+            );
+
+            // Revert to the snapshot and reset the agent's wallets.
+            chain.revert(id).await?;
+            alice.reset(Default::default());
+            bob.reset(Default::default());
+        }
+
+        Ok(())
+    }
+}