From c43faa6679ab2c97bb2a80146ef15ab6ab0add53 Mon Sep 17 00:00:00 2001
From: Christian Borst <christianoborst@gmail.com>
Date: Tue, 21 Jan 2025 12:14:52 -0500
Subject: [PATCH] Shorten upgrade test, make it more reliable

---
 .../test_runner/src/bootstrapping.rs          | 44 +++++++++++++------
 .../test_runner/src/tests/upgrade.rs          |  5 +--
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/integration_tests/test_runner/src/bootstrapping.rs b/integration_tests/test_runner/src/bootstrapping.rs
index 1ce04327..75fde575 100644
--- a/integration_tests/test_runner/src/bootstrapping.rs
+++ b/integration_tests/test_runner/src/bootstrapping.rs
@@ -14,13 +14,14 @@ use clarity::Uint256;
 use deep_space::private_key::CosmosPrivateKey;
 use deep_space::private_key::PrivateKey;
 use deep_space::Contact;
+use num::Zero;
 use std::fs::File;
 use std::io::{BufRead, BufReader, Read, Write};
 use std::os::fd::{FromRawFd, IntoRawFd};
 use std::path::Path;
 use std::process::{Command, ExitStatus, Stdio};
 use std::thread;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use web30::client::Web3;
 use web30::jsonrpc::error::Web3Error;
 
@@ -251,21 +252,36 @@ pub async fn send_erc20s_to_evm_users(
 
     // The users have been funded, skip sending erc20s
     info!("Checking for existing balances, might skip funding");
-    if !web3
-        .get_erc20_balance(
-            *erc20_contracts.first().unwrap(),
-            *destinations.first().unwrap(),
-        )
-        .await
-        .unwrap()
-        .is_zero()
-    {
-        return Ok(());
+
+    // There's a potential for the jsonrpc endpoint to be slow to start (especially in CI on the upgrade test)
+    // so we give the endpoint a bit to spin up before concluding that
+    let start = Instant::now();
+    let mut current_balance = Uint256::zero();
+    while start - Instant::now() < OPERATION_TIMEOUT {
+        match web3
+            .get_erc20_balance(
+                *erc20_contracts.first().unwrap(),
+                *destinations.first().unwrap(),
+            )
+            .await
+        {
+            Ok(balance) => {
+                current_balance = balance;
+                break;
+            }
+            Err(_) => {
+                continue;
+            }
+        }
     }
 
-    info!("Actually funding EVM users with the ERC20s");
-    for erc20 in erc20_contracts {
-        send_erc20_bulk(amount, erc20, &destinations, web3).await;
+    if current_balance.is_zero() {
+        info!("Actually funding EVM users with the ERC20s");
+        for erc20 in erc20_contracts {
+            send_erc20_bulk(amount, erc20, &destinations, web3).await;
+        }
+    } else {
+        info!("Skipping funding, balances already exist");
     }
     Ok(())
 }
diff --git a/integration_tests/test_runner/src/tests/upgrade.rs b/integration_tests/test_runner/src/tests/upgrade.rs
index 60c13094..41c51b81 100644
--- a/integration_tests/test_runner/src/tests/upgrade.rs
+++ b/integration_tests/test_runner/src/tests/upgrade.rs
@@ -1,9 +1,6 @@
 use crate::utils::{
     execute_upgrade_proposal, wait_for_block, UpgradeProposalParams, ValidatorKeys, EVM_USER_KEYS,
 };
-use althea_proto::cosmos_sdk_proto::cosmos::distribution::v1beta1::{
-    query_client::QueryClient as DistributionQueryClient, QueryParamsRequest,
-};
 use clarity::Address as EthAddress;
 use deep_space::client::ChainStatus;
 use deep_space::{Contact, CosmosPrivateKey};
@@ -16,7 +13,7 @@ use super::microtx_fees::microtx_fees_test;
 use super::native_token::native_token_test;
 
 pub const UPGRADE_NAME: &str = "example";
-const UPGRADE_BLOCK_DELTA: u64 = 20;
+const UPGRADE_BLOCK_DELTA: u64 = 30;
 
 /// Perform a series of integration tests to seed the system with data, then submit and pass a chain
 /// upgrade proposal