From e2c74dcdba1bcbcac7d6ef806262e12d19f9c0c6 Mon Sep 17 00:00:00 2001
From: ekexium <ekexium@gmail.com>
Date: Wed, 23 Sep 2020 14:24:39 +0800
Subject: [PATCH] fix integration tests

Signed-off-by: ekexium <ekexium@gmail.com>
---
 tests/integration_tests.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 9c7b412..f5e2669 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -35,11 +35,11 @@ async fn delete_all_raw() -> Fallible<()> {
 
 async fn delete_all_txn() -> Fallible<()> {
     let config = Config::new(pd_addrs());
-    let txn_client = TransactionClient::new(config).await?;
+    let txn_client = TransactionClient::new(config).await?.with_key_only(true);
     let mut txn = txn_client.begin().await?;
 
     loop {
-        let mut keys = txn.scan(vec![].., SCAN_BATCH_SIZE, true).await?.peekable();
+        let mut keys = txn.scan(vec![].., SCAN_BATCH_SIZE).await?.peekable();
 
         if keys.peek().is_none() {
             break;
@@ -177,6 +177,9 @@ async fn raw_bank_transfer() -> Fallible<()> {
         let mut alice_balance = get_u32(&client, alice.clone()).await?;
         let bob = chosen_people[1];
         let mut bob_balance = get_u32(&client, bob.clone()).await?;
+        if alice_balance == 0 {
+            continue;
+        }
         let transfer = rng.gen_range(0, alice_balance);
         alice_balance -= transfer;
         bob_balance += transfer;
@@ -211,7 +214,7 @@ async fn txn_bank_transfer() -> Fallible<()> {
     for person in &people {
         let init = rng.gen::<u8>() as u32;
         sum += init as u32;
-        txn.set(person.clone(), init.to_be_bytes().to_vec()).await?;
+        txn.put(person.clone(), init.to_be_bytes().to_vec()).await?;
     }
     txn.commit().await?;
 
@@ -223,12 +226,15 @@ async fn txn_bank_transfer() -> Fallible<()> {
         let mut alice_balance = get_txn_u32(&txn, alice.clone()).await?;
         let bob = chosen_people[1];
         let mut bob_balance = get_txn_u32(&txn, bob.clone()).await?;
+        if alice_balance == 0 {
+            continue;
+        }
         let transfer = rng.gen_range(0, alice_balance);
         alice_balance -= transfer;
         bob_balance += transfer;
-        txn.set(alice.clone(), alice_balance.to_be_bytes().to_vec())
+        txn.put(alice.clone(), alice_balance.to_be_bytes().to_vec())
             .await?;
-        txn.set(bob.clone(), bob_balance.to_be_bytes().to_vec())
+        txn.put(bob.clone(), bob_balance.to_be_bytes().to_vec())
             .await?;
         txn.commit().await?;
     }