From d672e86f0982e538e23451769a7077a572c97cb5 Mon Sep 17 00:00:00 2001
From: kitWarse <278602811+kitWarse@users.noreply.github.com>
Date: Sun, 31 May 2026 13:20:28 +0100
Subject: [PATCH] test: add sequential buy helper

---
 creator-keys/tests/buy_quote_monotonicity.rs | 31 +++++++++-----------
 creator-keys/tests/contract_test_env/mod.rs  | 19 ++++++++++++
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/creator-keys/tests/buy_quote_monotonicity.rs b/creator-keys/tests/buy_quote_monotonicity.rs
index b3f4e5e..2f24488 100644
--- a/creator-keys/tests/buy_quote_monotonicity.rs
+++ b/creator-keys/tests/buy_quote_monotonicity.rs
@@ -8,7 +8,8 @@
 mod contract_test_env;
 
 use contract_test_env::{
-    register_creator_keys, register_test_creator, set_pricing_and_fees, test_env_with_auths,
+    perform_incrementing_buys, register_creator_keys, register_test_creator, set_pricing_and_fees,
+    test_env_with_auths,
 };
 use creator_keys::CreatorKeysContractClient;
 use soroban_sdk::{testutils::Address as _, Address, Env};
@@ -20,18 +21,6 @@ fn setup_with_fees<'a>(env: &'a Env, price: i128) -> (CreatorKeysContractClient<
     (client, creator)
 }
 
-fn buy_n(
-    client: &CreatorKeysContractClient<'_>,
-    creator: &Address,
-    buyer: &Address,
-    n: u32,
-    price: i128,
-) {
-    for _ in 0..n {
-        client.buy_key(creator, buyer, &price);
-    }
-}
-
 // ── Monotonicity: fixed price means each quote is identical ───────────��───
 
 #[test]
@@ -70,9 +59,11 @@ fn test_buy_quote_price_unchanged_after_five_buys() {
     let buyer = Address::generate(&env);
 
     let before = client.get_buy_quote(&creator);
-    buy_n(&client, &creator, &buyer, 5, 500);
+    let state = perform_incrementing_buys(&client, &creator, &buyer, 5, 500, 1);
     let after = client.get_buy_quote(&creator);
 
+    assert_eq!(state.supply, 5);
+    assert_eq!(state.key_balance, 5);
     assert_eq!(before.price, after.price, "price must be deterministic");
     assert_eq!(before.total_amount, after.total_amount);
 }
@@ -103,9 +94,11 @@ fn test_buy_quote_total_amount_ordering_is_deterministic_small_range() {
     let buyer = Address::generate(&env);
 
     let q_start = client.get_buy_quote(&creator);
-    buy_n(&client, &creator, &buyer, 3, 1_000);
+    let state = perform_incrementing_buys(&client, &creator, &buyer, 3, 1_000, 1);
     let q_after = client.get_buy_quote(&creator);
 
+    assert_eq!(state.supply, 3);
+    assert_eq!(state.key_balance, 3);
     // Fixed price: total_amount should be unchanged.
     assert_eq!(q_start.total_amount, q_after.total_amount);
     // Fees must be non-negative.
@@ -119,9 +112,11 @@ fn test_buy_quote_fees_sum_to_total_minus_price() {
     let (client, creator) = setup_with_fees(&env, 1_000);
     let buyer = Address::generate(&env);
 
-    buy_n(&client, &creator, &buyer, 2, 1_000);
+    let state = perform_incrementing_buys(&client, &creator, &buyer, 2, 1_000, 1);
     let q = client.get_buy_quote(&creator);
 
+    assert_eq!(state.supply, 2);
+    assert_eq!(state.key_balance, 2);
     // total_amount = price + creator_fee + protocol_fee for a buy quote
     assert_eq!(
         q.total_amount,
@@ -165,8 +160,10 @@ fn test_buy_quote_total_amount_never_below_price() {
     let (client, creator) = setup_with_fees(&env, 10_000);
     let buyer = Address::generate(&env);
 
-    buy_n(&client, &creator, &buyer, 10, 10_000);
+    let state = perform_incrementing_buys(&client, &creator, &buyer, 10, 10_000, 1);
 
+    assert_eq!(state.supply, 10);
+    assert_eq!(state.key_balance, 10);
     let q = client.get_buy_quote(&creator);
     assert!(
         q.total_amount >= q.price,
diff --git a/creator-keys/tests/contract_test_env/mod.rs b/creator-keys/tests/contract_test_env/mod.rs
index 72c4a45..b7f3b5f 100644
--- a/creator-keys/tests/contract_test_env/mod.rs
+++ b/creator-keys/tests/contract_test_env/mod.rs
@@ -168,6 +168,25 @@ pub fn capture_snapshot(
     }
 }
 
+/// Performs `count` buys with payment amounts that increment from `starting_amount`.
+///
+/// Returns the resulting observable state for the creator and buyer after the sequence.
+pub fn perform_incrementing_buys(
+    client: &CreatorKeysContractClient<'_>,
+    creator: &Address,
+    buyer: &Address,
+    count: u32,
+    starting_amount: i128,
+    amount_step: i128,
+) -> ContractStateSnapshot {
+    for buy_index in 0..count {
+        let payment = starting_amount + i128::from(buy_index) * amount_step;
+        client.buy_key(creator, buyer, &payment);
+    }
+
+    capture_snapshot(client, creator, buyer)
+}
+
 impl ContractStateSnapshot {
     /// Asserts that `self` and `other` are identical, failing with a descriptive message if not.
     pub fn assert_unchanged(&self, after: &ContractStateSnapshot) {