feat: check balances equal deposits

vega_sim/api/data_raw.py

@@ -904,3 +904,19 @@ def list_stop_orders(
         request_func=lambda x: data_client.ListStopOrders(x).orders,
         extraction_func=lambda res: [i.node for i in res.edges],
     )
+
+
+@_retry(3)
+def list_deposits(
+    data_client: vac.trading_data_grpc_v2,
+    party_id: Optional[str] = None,
+    date_range: Optional[vega_protos.vega.DateRange] = None,
+) -> List[vega_protos.vega.Deposit]:
+    base_request = data_node_protos_v2.trading_data.ListDepositsRequest()
+    if party_id is not None:
+        setattr(base_request, "party_id", party_id)
+    return unroll_v2_pagination(
+        base_request=base_request,
+        request_func=lambda x: data_client.ListDeposits(x).deposits,
+        extraction_func=lambda res: [i.node for i in res.edges],
+    )

vega_sim/environment/environment.py

@@ -274,10 +274,10 @@ def _run(
                     f"Environment run at step {i}. Pausing to allow inspection of"
                     " state. Press Enter to continue"
                 )
-
             if step_end_callback is not None:
                 step_end_callback()
 
+        vega.check_balances_equal_deposits()
         logger.info(f"Run took {(datetime.datetime.now() - start).seconds}s")
 
         if pause_at_completion:

vega_sim/scenario/fuzzed_markets/run_fuzz_test.py

@@ -24,6 +24,7 @@ def _run(
     console: bool = False,
     output: bool = False,
     output_dir: str = "fuzz_plots",
+    lite: bool = False,
     core_metrics_port: int = 2723,
     data_node_metrics_port: int = 3651,
     perp_market_probability: float = 1.0,
@@ -35,6 +36,7 @@ def _run(
         transactions_per_block=4096,
         perps_market_probability=perp_market_probability,
         output=output,
+        lite=lite,
     )
 
     with VegaServiceNull(
@@ -98,6 +100,7 @@ def _run(
         action="store_true",
     )
     parser.add_argument("--console", action="store_true")
+    parser.add_argument("-l", "--lite", action="store_true")
     parser.add_argument("--core-metrics-port", default=2723, type=int)
     parser.add_argument("--data-node-metrics-port", default=3651, type=int)
     args = parser.parse_args()
@@ -111,6 +114,7 @@ def _run(
         steps=args.steps,
         console=args.console,
         output=True,
+        lite=args.lite,
         core_metrics_port=args.core_metrics_port,
         data_node_metrics_port=args.data_node_metrics_port,
     )

vega_sim/scenario/fuzzed_markets/scenario.py

@@ -153,6 +153,7 @@ def __init__(
         step_length_seconds: Optional[float] = None,
         fuzz_market_config: Optional[dict] = None,
         output: bool = True,
+        lite: bool = False,
     ):
         if perps_market_probability < 0 or perps_market_probability > 1:
             raise ValueError(
@@ -183,6 +184,7 @@ def __init__(
         self.transactions_per_block = transactions_per_block
 
         self.output = output
+        self.lite = lite
 
     def configure_agents(
         self,
@@ -215,7 +217,7 @@ def configure_agents(
             )
         )
 
-        for i_market in range(self.n_markets):
+        for i_market in range(self.n_markets if not self.lite else 1):
             # Determine if we should use perps market, otherwise use futures
             perps_market = self.random_state.random() < self.perps_probability
 
@@ -399,7 +401,7 @@ def configure_agents(
                         f"MARKET_{str(i_market).zfill(3)}_AGENT_{str(i_agent).zfill(3)}"
                     ),
                 )
-                for i_agent in range(10)
+                for i_agent in range(10 if not self.lite else 1)
             ]
 
             market_agents["risky_traders"] = [
@@ -415,7 +417,7 @@ def configure_agents(
                     tag=f"MARKET_{str(i_market).zfill(3)}_SIDE_{side}_AGENT_{str(i_agent).zfill(3)}",
                 )
                 for side in ["SIDE_BUY", "SIDE_SELL"]
-                for i_agent in range(10)
+                for i_agent in range(10 if not self.lite else 1)
             ]
 
             market_agents["risky_liquidity_providers"] = [
@@ -429,10 +431,10 @@ def configure_agents(
                     step_bias=0.1,
                     tag=f"HIGH_RISK_LPS_MARKET_{str(i_market).zfill(3)}_AGENT_{str(i_agent).zfill(3)}",
                 )
-                for i_agent in range(5)
+                for i_agent in range(5 if not self.lite else 1)
             ]
 
-            for i_agent in range(45):
+            for i_agent in range(45 if not self.lite else 1):
                 market_agents["risky_liquidity_providers"].append(
                     RiskySimpleLiquidityProvider(
                         wallet_name="risky_liquidity_providers",

vega_sim/service.py

@@ -82,6 +82,15 @@ class VegaFaucetError(Exception):
     pass
 
 
+class BalanceDepositInequity(Exception):
+    def __init__(self, asset, total_balance_amount, total_deposit_amount):
+        if total_balance_amount > total_deposit_amount:
+            msg = f"Balance in accounts greater than deposited funds ({total_balance_amount}>{total_deposit_amount}) for asset {asset}"
+        if total_balance_amount < total_deposit_amount:
+            msg = f"Balance in accounts less than deposited funds ({total_balance_amount}<{total_deposit_amount}) for asset {asset}"
+        super().__init__(msg)
+
+
 class MarketStateUpdateType(Enum):
     Unspecified = (
         vega_protos.governance.MarketStateUpdateType.MARKET_STATE_UPDATE_TYPE_UNSPECIFIED
@@ -3476,3 +3485,39 @@ def submit_proposal(
             self.wait_fn(int(time_to_enactment / self.seconds_per_block) + 1)
 
         self.wait_for_thread_catchup()
+
+    def check_balances_equal_deposits(self):
+        for attempts in range(100):
+            asset_balance_map = defaultdict(lambda: 0)
+            asset_deposit_map = defaultdict(lambda: 0)
+
+            for account in data_raw.list_accounts(
+                data_client=self.trading_data_client_v2
+            ):
+                asset_balance_map[account.asset] += int(account.balance)
+            for deposit in data_raw.list_deposits(
+                data_client=self.trading_data_client_v2
+            ):
+                asset_deposit_map[deposit.asset] += int(deposit.amount)
+
+            try:
+                for asset in asset_balance_map:
+                    total_balance_amount = asset_balance_map[asset]
+                    total_deposit_amount = asset_deposit_map[asset]
+                    assert asset_balance_map[asset] == asset_deposit_map[asset]
+                    logging.debug(
+                        f"Balance in accounts matches deposited funds for asset {asset}"
+                    )
+                return
+            except AssertionError:
+                logging.debug(
+                    "Balances don't match deposits, waiting to ensure datanode has finished consuming events."
+                )
+                time.sleep(0.0001 * 1.1**attempts)
+                continue
+
+        raise BalanceDepositInequity(
+            asset=asset,
+            total_balance_amount=total_balance_amount,
+            total_deposit_amount=total_deposit_amount,
+        )