worker/tests: Migrate to async database queries

Author: Turbo87

crates/crates_io_worker/tests/runner.rs

@@ -5,38 +5,41 @@ use crates_io_worker::{BackgroundJob, Runner};
 use diesel::prelude::*;
 use diesel_async::pooled_connection::deadpool::Pool;
 use diesel_async::pooled_connection::AsyncDieselConnectionManager;
-use diesel_async::AsyncPgConnection;
+use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl};
 use insta::assert_compact_json_snapshot;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 use tokio::sync::Barrier;
 
-fn all_jobs(conn: &mut PgConnection) -> Vec<(String, Value)> {
+async fn all_jobs(conn: &mut AsyncPgConnection) -> Vec<(String, Value)> {
     background_jobs::table
         .select((background_jobs::job_type, background_jobs::data))
         .get_results(conn)
+        .await
         .unwrap()
 }
 
-fn job_exists(id: i64, conn: &mut PgConnection) -> bool {
+async fn job_exists(id: i64, conn: &mut AsyncPgConnection) -> bool {
     background_jobs::table
         .find(id)
         .select(background_jobs::id)
         .get_result::<i64>(conn)
+        .await
         .optional()
         .unwrap()
         .is_some()
 }
 
-fn job_is_locked(id: i64, conn: &mut PgConnection) -> bool {
+async fn job_is_locked(id: i64, conn: &mut AsyncPgConnection) -> bool {
     background_jobs::table
         .find(id)
         .select(background_jobs::id)
         .for_update()
         .skip_locked()
         .get_result::<i64>(conn)
+        .await
         .optional()
         .unwrap()
         .is_none()
@@ -73,22 +76,25 @@ async fn jobs_are_locked_when_fetched() {
 
     let runner = runner(test_database.url(), test_context.clone()).register_job_type::<TestJob>();
 
-    let mut conn = test_database.connect();
-    let job_id = TestJob.enqueue(&mut conn).unwrap().unwrap();
+    let mut conn = AsyncPgConnection::establish(test_database.url())
+        .await
+        .unwrap();
+
+    let job_id = TestJob.async_enqueue(&mut conn).await.unwrap().unwrap();
 
-    assert!(job_exists(job_id, &mut conn));
-    assert!(!job_is_locked(job_id, &mut conn));
+    assert!(job_exists(job_id, &mut conn).await);
+    assert!(!job_is_locked(job_id, &mut conn).await);
 
     let runner = runner.start();
     test_context.job_started_barrier.wait().await;
 
-    assert!(job_exists(job_id, &mut conn));
-    assert!(job_is_locked(job_id, &mut conn));
+    assert!(job_exists(job_id, &mut conn).await);
+    assert!(job_is_locked(job_id, &mut conn).await);
 
     test_context.assertions_finished_barrier.wait().await;
     runner.wait_for_shutdown().await;
 
-    assert!(!job_exists(job_id, &mut conn));
+    assert!(!job_exists(job_id, &mut conn).await);
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -105,26 +111,30 @@ async fn jobs_are_deleted_when_successfully_run() {
         }
     }
 
-    fn remaining_jobs(conn: &mut PgConnection) -> i64 {
+    async fn remaining_jobs(conn: &mut AsyncPgConnection) -> i64 {
         background_jobs::table
             .count()
             .get_result(&mut *conn)
+            .await
             .unwrap()
     }
 
     let test_database = TestDatabase::new();
 
     let runner = runner(test_database.url(), ()).register_job_type::<TestJob>();
 
-    let mut conn = test_database.connect();
-    assert_eq!(remaining_jobs(&mut conn), 0);
+    let mut conn = AsyncPgConnection::establish(test_database.url())
+        .await
+        .unwrap();
+
+    assert_eq!(remaining_jobs(&mut conn).await, 0);
 
-    TestJob.enqueue(&mut conn).unwrap();
-    assert_eq!(remaining_jobs(&mut conn), 1);
+    TestJob.async_enqueue(&mut conn).await.unwrap();
+    assert_eq!(remaining_jobs(&mut conn).await, 1);
 
     let runner = runner.start();
     runner.wait_for_shutdown().await;
-    assert_eq!(remaining_jobs(&mut conn), 0);
+    assert_eq!(remaining_jobs(&mut conn).await, 0);
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -155,8 +165,11 @@ async fn failed_jobs_do_not_release_lock_before_updating_retry_time() {
 
     let runner = runner(test_database.url(), test_context.clone()).register_job_type::<TestJob>();
 
-    let mut conn = test_database.connect();
-    TestJob.enqueue(&mut conn).unwrap();
+    let mut conn = AsyncPgConnection::establish(test_database.url())
+        .await
+        .unwrap();
+
+    TestJob.async_enqueue(&mut conn).await.unwrap();
 
     let runner = runner.start();
     test_context.job_started_barrier.wait().await;
@@ -169,15 +182,17 @@ async fn failed_jobs_do_not_release_lock_before_updating_retry_time() {
         .select(background_jobs::id)
         .filter(background_jobs::retries.eq(0))
         .for_update()
-        .load::<i64>(&mut *conn)
+        .load::<i64>(&mut conn)
+        .await
         .unwrap();
     assert_eq!(available_jobs.len(), 0);
 
     // Sanity check to make sure the job actually is there
     let total_jobs_including_failed = background_jobs::table
         .select(background_jobs::id)
         .for_update()
-        .load::<i64>(&mut *conn)
+        .load::<i64>(&mut conn)
+        .await
         .unwrap();
     assert_eq!(total_jobs_including_failed.len(), 1);
 
@@ -202,9 +217,11 @@ async fn panicking_in_jobs_updates_retry_counter() {
 
     let runner = runner(test_database.url(), ()).register_job_type::<TestJob>();
 
-    let mut conn = test_database.connect();
+    let mut conn = AsyncPgConnection::establish(test_database.url())
+        .await
+        .unwrap();
 
-    let job_id = TestJob.enqueue(&mut conn).unwrap().unwrap();
+    let job_id = TestJob.async_enqueue(&mut conn).await.unwrap().unwrap();
 
     let runner = runner.start();
     runner.wait_for_shutdown().await;
@@ -213,7 +230,8 @@ async fn panicking_in_jobs_updates_retry_counter() {
         .find(job_id)
         .select(background_jobs::retries)
         .for_update()
-        .first::<i32>(&mut *conn)
+        .first::<i32>(&mut conn)
+        .await
         .unwrap();
     assert_eq!(tries, 1);
 }
@@ -264,33 +282,35 @@ async fn jobs_can_be_deduplicated() {
 
     let runner = runner(test_database.url(), test_context.clone()).register_job_type::<TestJob>();
 
-    let mut conn = test_database.connect();
+    let mut conn = AsyncPgConnection::establish(test_database.url())
+        .await
+        .unwrap();
 
     // Enqueue first job
-    assert_some!(TestJob::new("foo").enqueue(&mut conn).unwrap());
-    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}]]"#);
+    assert_some!(TestJob::new("foo").async_enqueue(&mut conn).await.unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn).await, @r#"[["test", {"value": "foo"}]]"#);
 
     // Try to enqueue the same job again, which should be deduplicated
-    assert_none!(TestJob::new("foo").enqueue(&mut conn).unwrap());
-    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}]]"#);
+    assert_none!(TestJob::new("foo").async_enqueue(&mut conn).await.unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn).await, @r#"[["test", {"value": "foo"}]]"#);
 
     // Start processing the first job
     let runner = runner.start();
     test_context.job_started_barrier.wait().await;
 
     // Enqueue the same job again, which should NOT be deduplicated,
     // since the first job already still running
-    assert_some!(TestJob::new("foo").enqueue(&mut conn).unwrap());
-    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
+    assert_some!(TestJob::new("foo").async_enqueue(&mut conn).await.unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn).await, @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
 
     // Try to enqueue the same job again, which should be deduplicated again
-    assert_none!(TestJob::new("foo").enqueue(&mut conn).unwrap());
-    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
+    assert_none!(TestJob::new("foo").async_enqueue(&mut conn).await.unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn).await, @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
 
     // Enqueue the same job but with different data, which should
     // NOT be deduplicated
-    assert_some!(TestJob::new("bar").enqueue(&mut conn).unwrap());
-    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}], ["test", {"value": "bar"}]]"#);
+    assert_some!(TestJob::new("bar").async_enqueue(&mut conn).await.unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn).await, @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}], ["test", {"value": "bar"}]]"#);
 
     // Resolve the final barrier to finish the test
     test_context.assertions_finished_barrier.wait().await;