Merge pull request #3 from Levana-Protocol/improved-tests

Improve the test suite
Levana-Protocol · Sep 3, 2024 · 1e96490 · 1e96490
2 parents d37acf3 + fd90f72
commit 1e96490
Show file tree

Hide file tree

Showing 6 changed files with 297 additions and 73 deletions.
diff --git a/contract/src/api.rs b/contract/src/api.rs
@@ -1,5 +1,3 @@
-use std::collections::BTreeMap;
-
 use crate::prelude::*;
 
 #[derive(Clone, Serialize, Deserialize, JsonSchema, Debug)]
@@ -101,7 +99,7 @@ pub struct OutcomeInfo {
 #[derive(Clone, Serialize, Deserialize, JsonSchema, Debug)]
 #[serde(rename_all = "snake_case")]
 pub struct PositionsResp {
-    pub outcomes: BTreeMap<OutcomeId, Token>,
+    pub outcomes: Vec<Token>,
 }
 
 #[derive(Serialize, Deserialize, JsonSchema, Debug)]

diff --git a/contract/src/execute.rs b/contract/src/execute.rs
@@ -158,14 +158,17 @@ fn deposit(
     let funds = deposit_amount.checked_sub(fee)?;
     let tokens = market.buy(outcome, funds)?;
     MARKETS.save(deps.storage, id, &market)?;
-    let mut share_info = SHARES
-        .may_load(deps.storage, (id, &info.sender))?
-        .unwrap_or_default();
-    *share_info
+    let mut share_info = ShareInfo::load(deps.storage, &market, &info.sender)?
+        .unwrap_or_else(|| ShareInfo::new(market.outcomes.len()));
+
+    assert_eq!(share_info.outcomes.len(), market.outcomes.len());
+
+    let outcome_tokens = share_info
         .outcomes
-        .entry(outcome)
-        .or_insert_with(Token::zero) += tokens;
-    SHARES.save(deps.storage, (id, &info.sender), &share_info)?;
+        .get_mut(usize::from(outcome.0) - 1)
+        .unwrap();
+    *outcome_tokens += tokens;
+    share_info.save(deps.storage, &market, &info.sender)?;
     Ok(Response::new().add_event(
         Event::new("deposit")
             .add_attribute("market-id", id.0.to_string())
@@ -194,13 +197,12 @@ fn withdraw(
         });
     }
 
-    let mut share_info = SHARES
-        .may_load(deps.storage, (id, &info.sender))?
+    let mut share_info = ShareInfo::load(deps.storage, &market, &info.sender)?
         .ok_or(Error::NoPositionsOnMarket { id })?;
 
     let user_tokens = share_info
         .outcomes
-        .get_mut(&outcome)
+        .get_mut(usize::from(outcome.0 - 1))
         .ok_or(Error::NoTokensFound { id, outcome })?;
 
     if *user_tokens < tokens {
@@ -212,7 +214,9 @@ fn withdraw(
         });
     }
 
-    SHARES.save(deps.storage, (id, &info.sender), &share_info)?;
+    *user_tokens -= tokens;
+
+    share_info.save(deps.storage, &market, &info.sender)?;
 
     let funds = market.sell(outcome, tokens)?;
 
@@ -248,7 +252,7 @@ fn set_winner(
 ) -> Result<Response> {
     let mut market = StoredMarket::load(deps.storage, id)?;
 
-    if env.block.time < market.withdrawal_stop_date {
+    if env.block.time < market.deposit_stop_date {
         return Err(Error::MarketStillActive {
             id,
             now: env.block.time,
@@ -291,21 +295,21 @@ fn set_winner(
 fn collect(deps: DepsMut, info: MessageInfo, id: MarketId) -> Result<Response> {
     let market = StoredMarket::load(deps.storage, id)?;
     let winner = market.winner.ok_or(Error::NoWinnerSet { id })?;
-    let mut share_info = SHARES
-        .may_load(deps.storage, (id, &info.sender))?
+    let mut share_info = ShareInfo::load(deps.storage, &market, &info.sender)?
         .ok_or(Error::NoPositionsOnMarket { id })?;
     if share_info.claimed_winnings {
         return Err(Error::AlreadyClaimedWinnings { id });
     }
     share_info.claimed_winnings = true;
-    let tokens = share_info
-        .outcomes
-        .get(&winner)
-        .ok_or(Error::NoTokensFound {
-            id,
-            outcome: winner,
-        })?;
-    SHARES.save(deps.storage, (id, &info.sender), &share_info)?;
+    let tokens =
+        share_info
+            .outcomes
+            .get(usize::from(winner.0 - 1))
+            .ok_or(Error::NoTokensFound {
+                id,
+                outcome: winner,
+            })?;
+    share_info.save(deps.storage, &market, &info.sender)?;
     let winnings = market.winnings_for(winner, *tokens)?;
 
     Ok(Response::new()

diff --git a/contract/src/query.rs b/contract/src/query.rs
@@ -25,9 +25,10 @@ fn market(deps: Deps, id: MarketId) -> Result<MarketResp> {
 
 fn positions(deps: Deps, id: MarketId, addr: String) -> Result<PositionsResp> {
     let addr = deps.api.addr_validate(&addr)?;
-    let outcomes = SHARES
-        .may_load(deps.storage, (id, &addr))?
-        .unwrap_or_default()
+    let market = StoredMarket::load(deps.storage, id)?;
+    let outcomes = ShareInfo::load(deps.storage, &market, &addr)?
+        .unwrap_or_else(|| ShareInfo::new(market.outcomes.len()))
         .outcomes;
+
     Ok(PositionsResp { outcomes })
 }
diff --git a/contract/src/state.rs b/contract/src/state.rs
@@ -1,5 +1,3 @@
-use std::collections::BTreeMap;
-
 use crate::prelude::*;
 
 pub const ADMIN: Item<Addr> = Item::new("admin");
@@ -10,7 +8,36 @@ pub const LAST_MARKET_ID: Item<MarketId> = Item::new("last-market-id");
 
 pub const MARKETS: Map<MarketId, StoredMarket> = Map::new("markets");
 
-pub const SHARES: Map<(MarketId, &Addr), ShareInfo> = Map::new("shares");
+const SHARES: Map<(MarketId, &Addr), ShareInfo> = Map::new("shares");
+
+impl ShareInfo {
+    pub fn load(
+        store: &dyn Storage,
+        market: &StoredMarket,
+        addr: &Addr,
+    ) -> Result<Option<ShareInfo>> {
+        match SHARES.may_load(store, (market.id, addr))? {
+            None => Ok(None),
+            Some(share_info) => {
+                if share_info.outcomes.len() != market.outcomes.len() {
+                    unreachable!()
+                } else {
+                    Ok(Some(share_info))
+                }
+            }
+        }
+    }
+
+    pub fn save(
+        &self,
+        store: &mut dyn Storage,
+        market: &StoredMarket,
+        addr: &Addr,
+    ) -> StdResult<()> {
+        assert_eq!(market.outcomes.len(), self.outcomes.len());
+        SHARES.save(store, (market.id, addr), self)
+    }
+}
 
 #[derive(Serialize, Deserialize)]
 pub struct StoredMarket {
@@ -58,8 +85,19 @@ pub struct StoredOutcome {
     pub total_tokens: Token,
 }
 
-#[derive(Serialize, Deserialize, Default)]
+#[derive(Serialize, Deserialize)]
 pub struct ShareInfo {
-    pub outcomes: BTreeMap<OutcomeId, Token>,
+    pub outcomes: Vec<Token>,
     pub claimed_winnings: bool,
 }
+
+impl ShareInfo {
+    pub fn new(outcome_count: usize) -> Self {
+        ShareInfo {
+            outcomes: std::iter::repeat(Token::zero())
+                .take(outcome_count)
+                .collect(),
+            claimed_winnings: false,
+        }
+    }
+}