From 0a6c7926ad7fcf2c519929515187729ed8cdbb81 Mon Sep 17 00:00:00 2001
From: Chris Czub <chris@penumbralabs.xyz>
Date: Mon, 29 Apr 2024 16:12:28 -0400
Subject: [PATCH] Add test for and fix arbitrage swap execution output amount

---
 .../core/component/dex/src/component/arb.rs   |  2 +-
 .../core/component/dex/src/component/tests.rs | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/crates/core/component/dex/src/component/arb.rs b/crates/core/component/dex/src/component/arb.rs
index b802ca2279..8017d348ca 100644
--- a/crates/core/component/dex/src/component/arb.rs
+++ b/crates/core/component/dex/src/component/arb.rs
@@ -114,7 +114,7 @@ pub trait Arbitrage: StateWrite + Sized {
                 amount: filled_input,
             },
             output: Value {
-                amount: arb_profit,
+                amount: filled_input + arb_profit,
                 asset_id: arb_token,
             },
         };
diff --git a/crates/core/component/dex/src/component/tests.rs b/crates/core/component/dex/src/component/tests.rs
index cfebe53716..75f16792b8 100644
--- a/crates/core/component/dex/src/component/tests.rs
+++ b/crates/core/component/dex/src/component/tests.rs
@@ -866,6 +866,9 @@ async fn basic_cycle_arb() -> anyhow::Result<()> {
 /// The issue was that we did not treat the spill price as a strict
 /// upper bound, which is necessary to ensure that the arbitrage logic
 /// terminates.
+///
+/// This test also ensures that the created `SwapExecution` has the
+/// correct data.
 async fn reproduce_arbitrage_loop_testnet_53() -> anyhow::Result<()> {
     let _ = tracing_subscriber::fmt::try_init();
     let storage = TempStorage::new().await?.apply_minimal_genesis().await?;
@@ -953,5 +956,46 @@ async fn reproduce_arbitrage_loop_testnet_53() -> anyhow::Result<()> {
     tracing::info!("fetching the `ArbExecution`");
     let arb_execution = state.arb_execution(0).await?.expect("arb was performed");
     tracing::info!(?arb_execution, "fetched arb execution!");
+
+    // Validate that the arb execution has the correct data:
+    // Validate the traces.
+    assert_eq!(
+        arb_execution.traces,
+        vec![
+            vec![
+                penumbra.value(1u32.into()),
+                test_usd.value(110u32.into()),
+                Value {
+                    amount: 1099999u64.into(),
+                    asset_id: penumbra.id()
+                }
+            ],
+            vec![
+                penumbra.value(1u32.into()),
+                test_usd.value(100u32.into()),
+                Value {
+                    amount: 999999u64.into(),
+                    asset_id: penumbra.id()
+                }
+            ]
+        ]
+    );
+
+    // Validate the input/output of the arb execution:
+    assert_eq!(
+        arb_execution.input,
+        Value {
+            amount: 2000000u64.into(),
+            asset_id: penumbra.id(),
+        }
+    );
+    assert_eq!(
+        arb_execution.output,
+        Value {
+            amount: 2099998u64.into(),
+            asset_id: penumbra.id(),
+        }
+    );
+
     Ok(())
 }