Skip to content

Commit 3fa4876

Browse files
Turbo87Turbo87
authored
test_db: Extract async_connect() fn (#10009)
1 parent b187729 commit 3fa4876

File tree

17 files changed

+49
-50
lines changed

17 files changed

+49
-50
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/crates_io_test_db/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ workspace = true
1010
[dependencies]
1111
crates_io_env_vars = { path = "../crates_io_env_vars" }
1212
diesel = { version = "=2.2.4", features = ["postgres", "r2d2"] }
13+
diesel-async = { version = "=0.5.1", features = ["postgres"] }
1314
diesel_migrations = { version = "=2.2.0", features = ["postgres"] }
1415
rand = "=0.8.5"
1516
tracing = "=0.1.40"

‎crates/crates_io_test_db/src/lib.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use crates_io_env_vars::required_var_parsed;
22
use diesel::prelude::*;
33
use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
44
use diesel::sql_query;
5+
use diesel_async::{AsyncConnection, AsyncPgConnection};
56
use diesel_migrations::{FileBasedMigrations, MigrationHarness};
67
use rand::Rng;
78
use std::sync::LazyLock;
@@ -128,6 +129,13 @@ impl TestDatabase {
128129
.get()
129130
.expect("Failed to get database connection")
130131
}
132+
133+
#[instrument(skip(self))]
134+
pub async fn async_connect(&self) -> AsyncPgConnection {
135+
AsyncPgConnection::establish(self.url())
136+
.await
137+
.expect("Failed to connect to database")
138+
}
131139
}
132140

133141
impl Drop for TestDatabase {

‎src/controllers/user/session.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -168,14 +168,13 @@ pub async fn logout(session: SessionExtension) -> Json<bool> {
168168
mod tests {
169169
use super::*;
170170
use crates_io_test_db::TestDatabase;
171-
use diesel_async::AsyncConnection;
172171

173172
#[tokio::test]
174173
async fn gh_user_with_invalid_email_doesnt_fail() {
175174
let emails = Emails::new_in_memory();
176175

177176
let test_db = TestDatabase::new();
178-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
177+
let mut conn = test_db.async_connect().await;
179178

180179
let gh_user = GithubUser {
181180
email: Some("String.Format(\"{0}.{1}@live.com\", FirstName, LastName)".into()),

‎src/index.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,14 +140,13 @@ mod tests {
140140
use crate::tests::builders::{CrateBuilder, VersionBuilder};
141141
use chrono::{Days, Utc};
142142
use crates_io_test_db::TestDatabase;
143-
use diesel_async::AsyncConnection;
144143
use insta::assert_json_snapshot;
145144

146145
#[tokio::test]
147146
async fn test_index_metadata() {
148147
let test_db = TestDatabase::new();
149148
let mut conn = test_db.connect();
150-
let mut async_conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
149+
let mut async_conn = test_db.async_connect().await;
151150

152151
let user_id = diesel::insert_into(users::table)
153152
.values((

‎src/models/category.rs‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -155,15 +155,14 @@ pub struct NewCategory<'a> {
155155
mod tests {
156156
use super::*;
157157
use crates_io_test_db::TestDatabase;
158-
use diesel_async::AsyncConnection;
159158
use diesel_async::RunQueryDsl;
160159

161160
#[tokio::test]
162161
async fn category_toplevel_excludes_subcategories() {
163162
use self::categories;
164163

165164
let test_db = TestDatabase::new();
166-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
165+
let mut conn = test_db.async_connect().await;
167166

168167
insert_into(categories::table)
169168
.values(&vec![
@@ -207,7 +206,7 @@ mod tests {
207206
};
208207

209208
let test_db = TestDatabase::new();
210-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
209+
let mut conn = test_db.async_connect().await;
211210

212211
insert_into(categories::table)
213212
.values(&vec![
@@ -238,7 +237,7 @@ mod tests {
238237
use self::categories;
239238

240239
let test_db = TestDatabase::new();
241-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
240+
let mut conn = test_db.async_connect().await;
242241

243242
insert_into(categories::table)
244243
.values(&vec![
@@ -287,7 +286,7 @@ mod tests {
287286
};
288287

289288
let test_db = TestDatabase::new();
290-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
289+
let mut conn = test_db.async_connect().await;
291290

292291
insert_into(categories::table)
293292
.values(&vec![
@@ -329,7 +328,7 @@ mod tests {
329328
};
330329

331330
let test_db = TestDatabase::new();
332-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
331+
let mut conn = test_db.async_connect().await;
333332

334333
insert_into(categories::table)
335334
.values(&vec![
@@ -376,7 +375,7 @@ mod tests {
376375
};
377376

378377
let test_db = TestDatabase::new();
379-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
378+
let mut conn = test_db.async_connect().await;
380379

381380
insert_into(categories::table)
382381
.values(&vec![

‎src/rate_limiter.rs‎

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -187,12 +187,11 @@ mod tests {
187187
use super::*;
188188
use crate::schema::users;
189189
use crates_io_test_db::TestDatabase;
190-
use diesel_async::AsyncConnection;
191190

192191
#[tokio::test]
193192
async fn default_rate_limits() -> anyhow::Result<()> {
194193
let test_db = TestDatabase::new();
195-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
194+
let mut conn = test_db.async_connect().await;
196195
let now = now();
197196

198197
// Set the defaults as if no env vars have been set in production
@@ -267,7 +266,7 @@ mod tests {
267266
#[tokio::test]
268267
async fn take_token_with_no_bucket_creates_new_one() -> anyhow::Result<()> {
269268
let test_db = TestDatabase::new();
270-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
269+
let mut conn = test_db.async_connect().await;
271270
let now = now();
272271

273272
let rate = SampleRateLimiter {
@@ -319,7 +318,7 @@ mod tests {
319318
#[tokio::test]
320319
async fn take_token_with_existing_bucket_modifies_existing_bucket() -> anyhow::Result<()> {
321320
let test_db = TestDatabase::new();
322-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
321+
let mut conn = test_db.async_connect().await;
323322
let now = now();
324323

325324
let rate = SampleRateLimiter {
@@ -345,7 +344,7 @@ mod tests {
345344
#[tokio::test]
346345
async fn take_token_after_delay_refills() -> anyhow::Result<()> {
347346
let test_db = TestDatabase::new();
348-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
347+
let mut conn = test_db.async_connect().await;
349348
let now = now();
350349

351350
let rate = SampleRateLimiter {
@@ -372,7 +371,7 @@ mod tests {
372371
#[tokio::test]
373372
async fn refill_subsecond_rate() -> anyhow::Result<()> {
374373
let test_db = TestDatabase::new();
375-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
374+
let mut conn = test_db.async_connect().await;
376375
// Subsecond rates have floating point rounding issues, so use a known
377376
// timestamp that rounds fine
378377
let now =
@@ -403,7 +402,7 @@ mod tests {
403402
#[tokio::test]
404403
async fn last_refill_always_advanced_by_multiple_of_rate() -> anyhow::Result<()> {
405404
let test_db = TestDatabase::new();
406-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
405+
let mut conn = test_db.async_connect().await;
407406
let now = now();
408407

409408
let rate = SampleRateLimiter {
@@ -435,7 +434,7 @@ mod tests {
435434
#[tokio::test]
436435
async fn zero_tokens_returned_when_user_has_no_tokens_left() -> anyhow::Result<()> {
437436
let test_db = TestDatabase::new();
438-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
437+
let mut conn = test_db.async_connect().await;
439438
let now = now();
440439

441440
let rate = SampleRateLimiter {
@@ -466,7 +465,7 @@ mod tests {
466465
#[tokio::test]
467466
async fn a_user_with_no_tokens_gets_a_token_after_exactly_rate() -> anyhow::Result<()> {
468467
let test_db = TestDatabase::new();
469-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
468+
let mut conn = test_db.async_connect().await;
470469
let now = now();
471470

472471
let rate = SampleRateLimiter {
@@ -494,7 +493,7 @@ mod tests {
494493
#[tokio::test]
495494
async fn tokens_never_refill_past_burst() -> anyhow::Result<()> {
496495
let test_db = TestDatabase::new();
497-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
496+
let mut conn = test_db.async_connect().await;
498497
let now = now();
499498

500499
let rate = SampleRateLimiter {
@@ -522,7 +521,7 @@ mod tests {
522521
#[tokio::test]
523522
async fn two_actions_dont_interfere_with_each_other() -> anyhow::Result<()> {
524523
let test_db = TestDatabase::new();
525-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
524+
let mut conn = test_db.async_connect().await;
526525
let now = now();
527526

528527
let mut config = HashMap::new();
@@ -571,7 +570,7 @@ mod tests {
571570
use diesel_async::RunQueryDsl;
572571

573572
let test_db = TestDatabase::new();
574-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
573+
let mut conn = test_db.async_connect().await;
575574
let now = now();
576575

577576
let rate = SampleRateLimiter {
@@ -609,7 +608,7 @@ mod tests {
609608
use diesel_async::RunQueryDsl;
610609

611610
let test_db = TestDatabase::new();
612-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
611+
let mut conn = test_db.async_connect().await;
613612
let now = now();
614613

615614
let rate = SampleRateLimiter {
@@ -669,7 +668,7 @@ mod tests {
669668
use diesel_async::RunQueryDsl;
670669

671670
let test_db = TestDatabase::new();
672-
let mut conn = AsyncPgConnection::establish(test_db.url()).await?;
671+
let mut conn = test_db.async_connect().await;
673672
let now = now();
674673
let user_id = new_user(&mut conn, "user").await?;
675674

‎src/tests/categories.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::schema::categories;
22
use crates_io_test_db::TestDatabase;
33
use diesel::*;
4-
use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl};
4+
use diesel_async::{AsyncPgConnection, RunQueryDsl};
55

66
const ALGORITHMS: &str = r#"
77
[algorithms]
@@ -50,7 +50,7 @@ async fn select_slugs(conn: &mut AsyncPgConnection) -> Vec<String> {
5050
#[tokio::test]
5151
async fn sync_adds_new_categories() {
5252
let test_db = TestDatabase::new();
53-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
53+
let mut conn = test_db.async_connect().await;
5454

5555
crate::boot::categories::sync_with_connection(ALGORITHMS_AND_SUCH, &mut conn)
5656
.await
@@ -63,7 +63,7 @@ async fn sync_adds_new_categories() {
6363
#[tokio::test]
6464
async fn sync_removes_missing_categories() {
6565
let test_db = TestDatabase::new();
66-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
66+
let mut conn = test_db.async_connect().await;
6767

6868
crate::boot::categories::sync_with_connection(ALGORITHMS_AND_SUCH, &mut conn)
6969
.await
@@ -79,7 +79,7 @@ async fn sync_removes_missing_categories() {
7979
#[tokio::test]
8080
async fn sync_adds_and_removes() {
8181
let test_db = TestDatabase::new();
82-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
82+
let mut conn = test_db.async_connect().await;
8383

8484
crate::boot::categories::sync_with_connection(ALGORITHMS_AND_SUCH, &mut conn)
8585
.await
@@ -95,7 +95,7 @@ async fn sync_adds_and_removes() {
9595
#[tokio::test]
9696
async fn test_real_categories() {
9797
let test_db = TestDatabase::new();
98-
let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
98+
let mut conn = test_db.async_connect().await;
9999

100100
const TOML: &str = include_str!("../boot/categories.toml");
101101
assert_ok!(crate::boot::categories::sync_with_connection(TOML, &mut conn).await);

‎src/tests/util/test_app.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use crates_io_test_db::TestDatabase;
2020
use crates_io_worker::Runner;
2121
use diesel::r2d2::{ConnectionManager, PooledConnection};
2222
use diesel::PgConnection;
23-
use diesel_async::{AsyncConnection, AsyncPgConnection};
23+
use diesel_async::AsyncPgConnection;
2424
use futures_util::TryStreamExt;
2525
use oauth2::{ClientId, ClientSecret};
2626
use regex::Regex;
@@ -119,8 +119,7 @@ impl TestApp {
119119

120120
/// Obtain an async database connection from the primary database pool.
121121
pub async fn async_db_conn(&self) -> AsyncPgConnection {
122-
let result = AsyncPgConnection::establish(self.0.test_database.url()).await;
123-
result.expect("Failed to get database connection")
122+
self.0.test_database.async_connect().await
124123
}
125124

126125
/// Create a new user with a verified email address in the database

‎src/typosquat/database.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,6 @@ mod tests {
176176
use super::*;
177177
use crate::typosquat::test_util::faker;
178178
use crates_io_test_db::TestDatabase;
179-
use diesel_async::AsyncConnection;
180179
use thiserror::Error;
181180

182181
#[tokio::test]
@@ -198,7 +197,7 @@ mod tests {
198197
faker::add_crate_to_team(&mut conn, &user_b, &top_b, &not_the_a_team)?;
199198
faker::add_crate_to_team(&mut conn, &user_b, &not_top_c, &not_the_a_team)?;
200199

201-
let mut async_conn = AsyncPgConnection::establish(test_db.url()).await?;
200+
let mut async_conn = test_db.async_connect().await;
202201
let top_crates = TopCrates::new(&mut async_conn, 2).await?;
203202

204203
// Let's ensure the top crates include what we expect (which is a and b, since we asked for

0 commit comments

Comments
 (0)