Skip to content

Commit 5d32527

Browse files
gusinaciogusinacio
committed
test: remove sleeps from all tests
Signed-off-by: Gustavo Inacio <[email protected]>
1 parent 7afff15 commit 5d32527

File tree

11 files changed

+125
-80
lines changed

11 files changed

+125
-80
lines changed

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/service/src/middleware/auth.rs‎

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,20 @@ pub use tap::tap_receipt_authorize;
1212
#[cfg(test)]
1313
mod tests {
1414
use std::sync::Arc;
15-
use std::time::Duration;
1615

1716
use axum::body::Body;
1817
use axum::http::{Request, Response};
1918
use reqwest::{header, StatusCode};
2019
use sqlx::PgPool;
2120
use tap_core::{manager::Manager, receipt::checks::CheckList};
22-
use tokio::time::sleep;
2321
use tower::{Service, ServiceBuilder, ServiceExt};
2422
use tower_http::auth::AsyncRequireAuthorizationLayer;
2523

2624
use crate::middleware::auth::{self, Bearer, OrExt};
2725
use crate::tap::IndexerTapContext;
28-
use test_assets::{create_signed_receipt, SignedReceiptRequest, TAP_EIP712_DOMAIN};
26+
use test_assets::{
27+
assert_while_retry, create_signed_receipt, SignedReceiptRequest, TAP_EIP712_DOMAIN,
28+
};
2929

3030
const BEARER_TOKEN: &str = "test";
3131

@@ -105,25 +105,13 @@ mod tests {
105105
assert_eq!(res.status(), StatusCode::OK);
106106

107107
// verify receipts
108-
if tokio::time::timeout(Duration::from_secs(1), async {
109-
loop {
110-
let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
111-
.fetch_all(&pgpool)
112-
.await
113-
.unwrap();
114-
115-
if result.is_empty() {
116-
sleep(Duration::from_millis(50)).await;
117-
} else {
118-
break;
119-
}
120-
}
121-
})
122-
.await
123-
.is_err()
124-
{
125-
panic!("Timeout assertion");
126-
}
108+
assert_while_retry!({
109+
let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
110+
.fetch_all(&pgpool)
111+
.await
112+
.unwrap();
113+
result.is_empty()
114+
});
127115
}
128116

129117
#[sqlx::test(migrations = "../../migrations")]

‎crates/service/src/middleware/auth/tap.rs‎

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,7 @@ mod tests {
7777

7878
use core::panic;
7979
use rstest::*;
80-
use std::{sync::Arc, time::Duration};
81-
use tokio::time::sleep;
80+
use std::sync::Arc;
8281
use tower::{Service, ServiceBuilder, ServiceExt};
8382

8483
use axum::{
@@ -96,7 +95,9 @@ mod tests {
9695
ReceiptWithState,
9796
},
9897
};
99-
use test_assets::{create_signed_receipt, SignedReceiptRequest, TAP_EIP712_DOMAIN};
98+
use test_assets::{
99+
assert_while_retry, create_signed_receipt, SignedReceiptRequest, TAP_EIP712_DOMAIN,
100+
};
100101
use tower_http::auth::AsyncRequireAuthorizationLayer;
101102

102103
use crate::{
@@ -181,25 +182,13 @@ mod tests {
181182
assert_eq!(res.status(), StatusCode::OK);
182183

183184
// verify receipts
184-
if tokio::time::timeout(Duration::from_secs(1), async {
185-
loop {
186-
let result = sqlx::query!("SELECT * FROM scalar_tap_receipts")
187-
.fetch_all(&pgpool)
188-
.await
189-
.unwrap();
190-
191-
if result.is_empty() {
192-
sleep(Duration::from_millis(50)).await;
193-
} else {
194-
break;
195-
}
196-
}
185+
assert_while_retry!({
186+
sqlx::query!("SELECT * FROM scalar_tap_receipts")
187+
.fetch_all(&pgpool)
188+
.await
189+
.unwrap()
190+
.is_empty()
197191
})
198-
.await
199-
.is_err()
200-
{
201-
panic!("Timeout assertion");
202-
}
203192
}
204193

205194
#[rstest]

‎crates/service/src/tap/checks/deny_list_check.rs‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@ pub struct DenyListCheck {
2020
sender_denylist: Arc<RwLock<HashSet<Address>>>,
2121
_sender_denylist_watcher_handle: Arc<tokio::task::JoinHandle<()>>,
2222
sender_denylist_watcher_cancel_token: tokio_util::sync::CancellationToken,
23+
24+
#[cfg(test)]
25+
notify: std::sync::Arc<tokio::sync::Notify>,
2326
}
2427

2528
impl DenyListCheck {
@@ -41,17 +44,24 @@ impl DenyListCheck {
4144
.await
4245
.expect("should be able to fetch the sender_denylist from the DB on startup");
4346

47+
#[cfg(test)]
48+
let notify = std::sync::Arc::new(tokio::sync::Notify::new());
49+
4450
let sender_denylist_watcher_cancel_token = tokio_util::sync::CancellationToken::new();
4551
let sender_denylist_watcher_handle = Arc::new(tokio::spawn(Self::sender_denylist_watcher(
4652
pgpool.clone(),
4753
pglistener,
4854
sender_denylist.clone(),
4955
sender_denylist_watcher_cancel_token.clone(),
56+
#[cfg(test)]
57+
notify.clone(),
5058
)));
5159
Self {
5260
sender_denylist,
5361
_sender_denylist_watcher_handle: sender_denylist_watcher_handle,
5462
sender_denylist_watcher_cancel_token,
63+
#[cfg(test)]
64+
notify,
5565
}
5666
}
5767

@@ -81,6 +91,7 @@ impl DenyListCheck {
8191
mut pglistener: PgListener,
8292
denylist: Arc<RwLock<HashSet<Address>>>,
8393
cancel_token: tokio_util::sync::CancellationToken,
94+
#[cfg(test)] notify: std::sync::Arc<tokio::sync::Notify>,
8495
) {
8596
#[derive(serde::Deserialize)]
8697
struct DenylistNotification {
@@ -132,6 +143,8 @@ impl DenyListCheck {
132143
.expect("should be able to reload the sender denylist")
133144
}
134145
}
146+
#[cfg(test)]
147+
notify.notify_one();
135148
}
136149
}
137150
}
@@ -246,8 +259,9 @@ mod tests {
246259
.await
247260
.unwrap();
248261

262+
deny_list_check.notify.notified().await;
263+
249264
// Check that the receipt is rejected
250-
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
251265
assert!(deny_list_check
252266
.check(&ctx, &checking_receipt)
253267
.await
@@ -265,11 +279,9 @@ mod tests {
265279
.await
266280
.unwrap();
267281

282+
deny_list_check.notify.notified().await;
283+
268284
// Check that the receipt is valid again
269-
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
270-
deny_list_check
271-
.check(&ctx, &checking_receipt)
272-
.await
273-
.unwrap();
285+
assert!(deny_list_check.check(&ctx, &checking_receipt).await.is_ok());
274286
}
275287
}

‎crates/service/src/tap/checks/value_check.rs‎

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ pub struct MinimumValue {
5353
watcher_cancel_token: tokio_util::sync::CancellationToken,
5454
updated_at: GracePeriod,
5555
grace_period: Duration,
56+
57+
#[cfg(test)]
58+
notify: std::sync::Arc<tokio::sync::Notify>,
5659
}
5760

5861
struct CostModelWatcher {
@@ -61,6 +64,9 @@ struct CostModelWatcher {
6164
cost_models: CostModelMap,
6265
global_model: GlobalModel,
6366
updated_at: GracePeriod,
67+
68+
#[cfg(test)]
69+
notify: std::sync::Arc<tokio::sync::Notify>,
6470
}
6571

6672
impl CostModelWatcher {
@@ -71,12 +77,15 @@ impl CostModelWatcher {
7177
global_model: GlobalModel,
7278
cancel_token: tokio_util::sync::CancellationToken,
7379
grace_period: GracePeriod,
80+
#[cfg(test)] notify: std::sync::Arc<tokio::sync::Notify>,
7481
) {
7582
let cost_model_watcher = CostModelWatcher {
7683
pgpool,
7784
global_model,
7885
cost_models,
7986
updated_at: grace_period,
87+
#[cfg(test)]
88+
notify,
8089
};
8190

8291
loop {
@@ -109,6 +118,8 @@ impl CostModelWatcher {
109118
// model cache.
110119
Err(_) => self.handle_unexpected_notification(payload).await,
111120
}
121+
#[cfg(test)]
122+
self.notify.notify_one();
112123
}
113124

114125
fn handle_insert(&self, deployment: String, model: String, variables: String) {
@@ -200,6 +211,9 @@ impl MinimumValue {
200211
'cost_models_update_notification'",
201212
);
202213

214+
#[cfg(test)]
215+
let notify = std::sync::Arc::new(tokio::sync::Notify::new());
216+
203217
let watcher_cancel_token = tokio_util::sync::CancellationToken::new();
204218
tokio::spawn(CostModelWatcher::cost_models_watcher(
205219
pgpool.clone(),
@@ -208,13 +222,17 @@ impl MinimumValue {
208222
global_model.clone(),
209223
watcher_cancel_token.clone(),
210224
updated_at.clone(),
225+
#[cfg(test)]
226+
notify.clone(),
211227
));
212228
Self {
213229
global_model,
214230
cost_model_map,
215231
watcher_cancel_token,
216232
updated_at,
217233
grace_period,
234+
#[cfg(test)]
235+
notify,
218236
}
219237
}
220238

@@ -347,7 +365,7 @@ enum CostModelNotification {
347365
#[cfg(test)]
348366
mod tests {
349367
use std::time::Duration;
350-
use test_assets::{create_signed_receipt, SignedReceiptRequest};
368+
use test_assets::{create_signed_receipt, flush_messages, SignedReceiptRequest};
351369

352370
use sqlx::PgPool;
353371
use tap_core::receipt::{checks::Check, Context, ReceiptWithState};
@@ -388,7 +406,8 @@ mod tests {
388406
// insert 2 cost models for different deployment_id
389407
let test_models = test::test_data();
390408
add_cost_models(&pgpool, to_db_models(test_models.clone())).await;
391-
sleep(Duration::from_millis(200)).await;
409+
410+
flush_messages(&check.notify).await;
392411

393412
assert_eq!(
394413
check.cost_model_map.read().unwrap().len(),
@@ -411,7 +430,7 @@ mod tests {
411430
.await
412431
.unwrap();
413432

414-
sleep(Duration::from_millis(200)).await;
433+
check.notify.notified().await;
415434

416435
assert_eq!(check.cost_model_map.read().unwrap().len(), 0);
417436
}
@@ -431,7 +450,8 @@ mod tests {
431450

432451
let global_model = global_cost_model();
433452
add_cost_models(&pgpool, vec![global_model.clone()]).await;
434-
sleep(Duration::from_millis(10)).await;
453+
454+
check.notify.notified().await;
435455

436456
assert!(check.global_model.read().unwrap().is_some());
437457
}
@@ -449,7 +469,7 @@ mod tests {
449469
.await
450470
.unwrap();
451471

452-
sleep(Duration::from_millis(10)).await;
472+
check.notify.notified().await;
453473

454474
assert_eq!(check.cost_model_map.read().unwrap().len(), 0);
455475
}
@@ -461,7 +481,9 @@ mod tests {
461481

462482
add_cost_models(&pgpool, to_db_models(test_models.clone())).await;
463483

464-
let check = MinimumValue::new(pgpool, Duration::from_secs(1)).await;
484+
let grace_period = Duration::from_secs(1);
485+
486+
let check = MinimumValue::new(pgpool, grace_period).await;
465487

466488
let deployment_id = test_models[0].deployment;
467489
let mut ctx = Context::new();
@@ -511,7 +533,7 @@ mod tests {
511533
"Should accept since its inside grace period "
512534
);
513535

514-
sleep(Duration::from_millis(1010)).await;
536+
sleep(grace_period + Duration::from_millis(10)).await;
515537

516538
assert!(
517539
check.check(&ctx, &receipt).await.is_err(),

‎crates/tap-agent/src/agent/sender_account.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,9 +1032,7 @@ pub mod tests {
10321032
use crate::agent::sender_account::ReceiptFees;
10331033
use crate::agent::sender_allocation::SenderAllocationMessage;
10341034
use crate::agent::unaggregated_receipts::UnaggregatedReceipts;
1035-
use crate::test::actors::{
1036-
create_mock_sender_allocation, flush_messages, MockSenderAllocation, TestableActor,
1037-
};
1035+
use crate::test::actors::{create_mock_sender_allocation, MockSenderAllocation, TestableActor};
10381036
use crate::test::{create_rav, store_rav_with_options, INDEXER, TAP_EIP712_DOMAIN_SEPARATOR};
10391037
use crate::{assert_not_triggered, assert_triggered};
10401038

@@ -1050,7 +1048,8 @@ pub mod tests {
10501048
use std::sync::Arc;
10511049
use std::time::{Duration, SystemTime, UNIX_EPOCH};
10521050
use test_assets::{
1053-
ALLOCATION_ID_0, ALLOCATION_ID_1, TAP_SENDER as SENDER, TAP_SIGNER as SIGNER,
1051+
flush_messages, ALLOCATION_ID_0, ALLOCATION_ID_1, TAP_SENDER as SENDER,
1052+
TAP_SIGNER as SIGNER,
10541053
};
10551054
use tokio::sync::watch::{self, Sender};
10561055
use tokio::sync::Notify;

‎crates/tap-agent/src/agent/sender_accounts_manager.rs‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -587,9 +587,7 @@ mod tests {
587587
use crate::agent::sender_account::tests::PREFIX_ID;
588588
use crate::agent::sender_account::{SenderAccountConfig, SenderAccountMessage};
589589
use crate::agent::sender_accounts_manager::{handle_notification, NewReceiptNotification};
590-
use crate::test::actors::{
591-
flush_messages, DummyActor, MockSenderAccount, MockSenderAllocation, TestableActor,
592-
};
590+
use crate::test::actors::{DummyActor, MockSenderAccount, MockSenderAllocation, TestableActor};
593591
use crate::test::{
594592
create_rav, create_received_receipt, store_rav, store_receipt, ALLOCATION_ID_0,
595593
ALLOCATION_ID_1, INDEXER, SENDER_2, SIGNER, TAP_EIP712_DOMAIN_SEPARATOR,
@@ -605,7 +603,7 @@ mod tests {
605603
use std::collections::{HashMap, HashSet};
606604
use std::sync::Arc;
607605
use std::time::Duration;
608-
use test_assets::TAP_SENDER as SENDER;
606+
use test_assets::{flush_messages, TAP_SENDER as SENDER};
609607
use tokio::sync::mpsc::error::TryRecvError;
610608
use tokio::sync::{mpsc, watch, Notify};
611609

‎crates/tap-agent/src/agent/sender_allocation.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -867,7 +867,7 @@ pub mod tests {
867867
unaggregated_receipts::UnaggregatedReceipts,
868868
},
869869
test::{
870-
actors::{create_mock_sender_account, flush_messages, TestableActor},
870+
actors::{create_mock_sender_account, TestableActor},
871871
create_rav, create_received_receipt, store_invalid_receipt, store_rav, store_receipt,
872872
INDEXER,
873873
},
@@ -892,8 +892,8 @@ pub mod tests {
892892
Context, ReceiptWithState,
893893
};
894894
use test_assets::{
895-
ALLOCATION_ID_0, TAP_EIP712_DOMAIN as TAP_EIP712_DOMAIN_SEPARATOR, TAP_SENDER as SENDER,
896-
TAP_SIGNER as SIGNER,
895+
flush_messages, ALLOCATION_ID_0, TAP_EIP712_DOMAIN as TAP_EIP712_DOMAIN_SEPARATOR,
896+
TAP_SENDER as SENDER, TAP_SIGNER as SIGNER,
897897
};
898898
use tokio::sync::{watch, Notify};
899899
use wiremock::{

0 commit comments

Comments
 (0)