Replace sleep with explicit wait for authenticator state tx (MystenLa…

…bs#14009)

crates/sui-e2e-tests/tests/zklogin_tests.rs

@@ -3,7 +3,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::BTreeSet;
-use tokio::time::{sleep, Duration};
 
 use sui_test_transaction_builder::TestTransactionBuilder;
 use sui_types::error::{SuiError, SuiResult};
@@ -86,16 +85,13 @@ async fn zklogin_end_to_end_test_with_auth_state_creation() {
     // Now wait until the next epoch, when the auth state object is created.
     test_cluster.wait_for_epoch(None).await;
 
-    // Wait for JWKs to be fetched and sequenced.
-    sleep(Duration::from_secs(10)).await;
-
     // run zklogin end to end test
     run_zklogin_end_to_end_test(test_cluster).await;
 }
 
 async fn run_zklogin_end_to_end_test(mut test_cluster: TestCluster) {
     // wait for JWKs to be fetched and sequenced.
-    sleep(Duration::from_secs(15)).await;
+    test_cluster.wait_for_authenticator_state_update().await;
 
     let rgp = test_cluster.get_reference_gas_price().await;
     let sender = test_cluster.get_address_0();
@@ -191,6 +187,7 @@ async fn test_conflicting_jwks() {
     use sui_json_rpc_types::TransactionFilter;
     use sui_types::base_types::ObjectID;
     use sui_types::transaction::{TransactionDataAPI, TransactionKind};
+    use tokio::time::Duration;
 
     let test_cluster = TestClusterBuilder::new()
         .with_epoch_duration_ms(45000)

crates/test-cluster/src/lib.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use futures::future::join_all;
+use futures::{future::join_all, StreamExt};
 use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
 use jsonrpsee::ws_client::WsClient;
 use jsonrpsee::ws_client::WsClientBuilder;
@@ -17,7 +17,9 @@ use sui_config::{Config, SUI_CLIENT_CONFIG, SUI_NETWORK_CONFIG};
 use sui_config::{NodeConfig, PersistedConfig, SUI_KEYSTORE_FILENAME};
 use sui_core::authority_aggregator::AuthorityAggregator;
 use sui_core::authority_client::NetworkAuthorityClient;
-use sui_json_rpc_types::SuiTransactionBlockResponse;
+use sui_json_rpc_types::{
+    SuiTransactionBlockEffectsAPI, SuiTransactionBlockResponse, TransactionFilter,
+};
 use sui_keys::keystore::{AccountKeystore, FileBasedKeystore, Keystore};
 use sui_node::SuiNodeHandle;
 use sui_protocol_config::{ProtocolVersion, SupportedProtocolVersions};
@@ -46,7 +48,7 @@ use sui_types::object::Object;
 use sui_types::sui_system_state::epoch_start_sui_system_state::EpochStartSystemStateTrait;
 use sui_types::sui_system_state::SuiSystemState;
 use sui_types::sui_system_state::SuiSystemStateTrait;
-use sui_types::transaction::{Transaction, TransactionData};
+use sui_types::transaction::{Transaction, TransactionData, TransactionDataAPI, TransactionKind};
 use tokio::time::{timeout, Instant};
 use tokio::{task::JoinHandle, time::sleep};
 use tracing::info;
@@ -404,6 +406,34 @@ impl TestCluster {
         }
     }
 
+    pub async fn wait_for_authenticator_state_update(&self) {
+        timeout(
+            Duration::from_secs(60),
+            self.fullnode_handle.sui_node.with_async(|node| async move {
+                let mut txns = node.state().subscription_handler.subscribe_transactions(
+                    TransactionFilter::ChangedObject(ObjectID::from_hex_literal("0x7").unwrap()),
+                );
+                let state = node.state();
+
+                while let Some(tx) = txns.next().await {
+                    let digest = *tx.transaction_digest();
+                    let tx = state
+                        .database
+                        .get_transaction_block(&digest)
+                        .unwrap()
+                        .unwrap();
+                    match &tx.data().intent_message().value.kind() {
+                        TransactionKind::EndOfEpochTransaction(_) => (),
+                        TransactionKind::AuthenticatorStateUpdate(_) => break,
+                        _ => panic!("{:?}", tx),
+                    }
+                }
+            }),
+        )
+        .await
+        .expect("Timed out waiting for authenticator state update");
+    }
+
     pub async fn test_transaction_builder(&self) -> TestTransactionBuilder {
         let (sender, gas) = self.wallet.get_one_gas_object().await.unwrap().unwrap();
         let rgp = self.get_reference_gas_price().await;