From 3d1adab81aea1d92de97b220755827548e653f7d Mon Sep 17 00:00:00 2001
From: Jort <jordan.jennings@mystenlabs.com>
Date: Mon, 27 Jan 2025 17:59:07 -0800
Subject: [PATCH] allow gas and sponsor in dry run

---
 .../dev_inspect/reference_return_value.move   |  1 +
 .../tests/dry_run/reference_return_value.move |  1 +
 .../src/test_adapter.rs                       | 36 +++++++++++--------
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/crates/sui-adapter-transactional-tests/tests/dev_inspect/reference_return_value.move b/crates/sui-adapter-transactional-tests/tests/dev_inspect/reference_return_value.move
index 24b1298379320..ce3a4cb2fdb52 100644
--- a/crates/sui-adapter-transactional-tests/tests/dev_inspect/reference_return_value.move
+++ b/crates/sui-adapter-transactional-tests/tests/dev_inspect/reference_return_value.move
@@ -14,4 +14,5 @@ module test::m {
 }
 
 //# programmable --sender A --inputs 0 --dev-inspect
+
 //> 0: test::m::return_ref(Input(0));
\ No newline at end of file
diff --git a/crates/sui-adapter-transactional-tests/tests/dry_run/reference_return_value.move b/crates/sui-adapter-transactional-tests/tests/dry_run/reference_return_value.move
index c8a749a2bac3d..5bfc5b6e13346 100644
--- a/crates/sui-adapter-transactional-tests/tests/dry_run/reference_return_value.move
+++ b/crates/sui-adapter-transactional-tests/tests/dry_run/reference_return_value.move
@@ -14,4 +14,5 @@ module test::m {
 }
 
 //# programmable --sender A --inputs 0 --dry-run
+
 //> 0: test::m::return_ref(Input(0));
\ No newline at end of file
diff --git a/crates/sui-transactional-test-runner/src/test_adapter.rs b/crates/sui-transactional-test-runner/src/test_adapter.rs
index 834b927033d45..d344654f88bc1 100644
--- a/crates/sui-transactional-test-runner/src/test_adapter.rs
+++ b/crates/sui-transactional-test-runner/src/test_adapter.rs
@@ -904,10 +904,14 @@ impl<'a> MoveTestAdapter<'a> for SuiTestAdapter {
                 } else if dry_run {
                     let gas_budget = gas_budget.unwrap_or(DEFAULT_GAS_BUDGET);
                     let gas_price = gas_price.unwrap_or(self.gas_price);
-                    let sender_address = self.get_sender(sender).address;
+                    let sender = self.get_sender(sender);
+                    let sponsor = sponsor.map_or(sender, |a| self.get_sender(Some(a)));
+
+                    let payment = self.get_payment(sponsor, gas_payment);
+
                     let transaction = TransactionData::new_programmable(
-                        sender_address,
-                        vec![],
+                        sender.address,
+                        vec![payment],
                         ProgrammableTransaction { inputs, commands },
                         gas_budget,
                         gas_price,
@@ -1401,6 +1405,7 @@ impl<'a> SuiTestAdapter {
         gas_budget: Option<u64>,
         policy: u8,
         gas_price: u64,
+        // dry_run: bool,
     ) -> anyhow::Result<Option<String>> {
         let modules_bytes = modules
             .iter()
@@ -1494,6 +1499,19 @@ impl<'a> SuiTestAdapter {
         })
     }
 
+    fn get_payment(&self, sponsor: &TestAccount, payment: Option<FakeID>) -> ObjectRef {
+        let payment = if let Some(payment) = payment {
+            self.fake_to_real_object_id(payment)
+                .expect("Could not find specified payment object")
+        } else {
+            sponsor.gas
+        };
+
+        self.get_object(&payment, None)
+            .unwrap()
+            .compute_object_reference()
+    }
+
     fn sign_sponsor_txn(
         &self,
         sender: Option<String>,
@@ -1508,17 +1526,7 @@ impl<'a> SuiTestAdapter {
         let sender = self.get_sender(sender);
         let sponsor = sponsor.map_or(sender, |a| self.get_sender(Some(a)));
 
-        let payment = if let Some(payment) = payment {
-            self.fake_to_real_object_id(payment)
-                .expect("Could not find specified payment object")
-        } else {
-            sponsor.gas
-        };
-
-        let payment_ref = self
-            .get_object(&payment, None)
-            .unwrap()
-            .compute_object_reference();
+        let payment_ref = self.get_payment(sender, payment);
 
         let data = txn_data(sender.address, sponsor.address, payment_ref);
         if sender.address == sponsor.address {