sqlx_core/testing/
mod.rs

1use std::future::Future;
2use std::time::Duration;
3
4use futures_core::future::BoxFuture;
5
6use base64::{engine::general_purpose::URL_SAFE, Engine as _};
7pub use fixtures::FixtureSnapshot;
8use sha2::{Digest, Sha512};
9
10use crate::connection::{ConnectOptions, Connection};
11use crate::database::Database;
12use crate::error::Error;
13use crate::executor::Executor;
14use crate::migrate::{Migrate, Migrator};
15use crate::pool::{Pool, PoolConnection, PoolOptions};
16
17mod fixtures;
18
19pub trait TestSupport: Database {
20    /// Get parameters to construct a `Pool` suitable for testing.
21    ///
22    /// This `Pool` instance will behave somewhat specially:
23    /// * all handles share a single global semaphore to avoid exceeding the connection limit
24    ///   on the database server.
25    /// * each invocation results in a different temporary database.
26    ///
27    /// The implementation may require `DATABASE_URL` to be set in order to manage databases.
28    /// The user credentials it contains must have the privilege to create and drop databases.
29    fn test_context(args: &TestArgs) -> BoxFuture<'_, Result<TestContext<Self>, Error>>;
30
31    fn cleanup_test(db_name: &str) -> BoxFuture<'_, Result<(), Error>>;
32
33    /// Cleanup any test databases that are no longer in-use.
34    ///
35    /// Returns a count of the databases deleted, if possible.
36    ///
37    /// The implementation may require `DATABASE_URL` to be set in order to manage databases.
38    /// The user credentials it contains must have the privilege to create and drop databases.
39    fn cleanup_test_dbs() -> BoxFuture<'static, Result<Option<usize>, Error>>;
40
41    /// Take a snapshot of the current state of the database (data only).
42    ///
43    /// This snapshot can then be used to generate test fixtures.
44    fn snapshot(conn: &mut Self::Connection)
45        -> BoxFuture<'_, Result<FixtureSnapshot<Self>, Error>>;
46
47    /// Generate a unique database name for the given test path.
48    fn db_name(args: &TestArgs) -> String {
49        let mut hasher = Sha512::new();
50        hasher.update(args.test_path.as_bytes());
51        let hash = hasher.finalize();
52        let hash = URL_SAFE.encode(&hash[..39]);
53        let db_name = format!("_sqlx_test_{}", hash).replace('-', "_");
54        debug_assert!(db_name.len() == 63);
55        db_name
56    }
57}
58
59pub struct TestFixture {
60    pub path: &'static str,
61    pub contents: &'static str,
62}
63
64pub struct TestArgs {
65    pub test_path: &'static str,
66    pub migrator: Option<&'static Migrator>,
67    pub fixtures: &'static [TestFixture],
68}
69
70pub trait TestFn {
71    type Output;
72
73    fn run_test(self, args: TestArgs) -> Self::Output;
74}
75
76pub trait TestTermination {
77    fn is_success(&self) -> bool;
78}
79
80pub struct TestContext<DB: Database> {
81    pub pool_opts: PoolOptions<DB>,
82    pub connect_opts: <DB::Connection as Connection>::Options,
83    pub db_name: String,
84}
85
86impl<DB, Fut> TestFn for fn(Pool<DB>) -> Fut
87where
88    DB: TestSupport + Database,
89    DB::Connection: Migrate,
90    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
91    Fut: Future,
92    Fut::Output: TestTermination,
93{
94    type Output = Fut::Output;
95
96    fn run_test(self, args: TestArgs) -> Self::Output {
97        run_test_with_pool(args, self)
98    }
99}
100
101impl<DB, Fut> TestFn for fn(PoolConnection<DB>) -> Fut
102where
103    DB: TestSupport + Database,
104    DB::Connection: Migrate,
105    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
106    Fut: Future,
107    Fut::Output: TestTermination,
108{
109    type Output = Fut::Output;
110
111    fn run_test(self, args: TestArgs) -> Self::Output {
112        run_test_with_pool(args, |pool| async move {
113            let conn = pool
114                .acquire()
115                .await
116                .expect("failed to acquire test pool connection");
117            let res = (self)(conn).await;
118            pool.close().await;
119            res
120        })
121    }
122}
123
124impl<DB, Fut> TestFn for fn(PoolOptions<DB>, <DB::Connection as Connection>::Options) -> Fut
125where
126    DB: Database + TestSupport,
127    DB::Connection: Migrate,
128    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
129    Fut: Future,
130    Fut::Output: TestTermination,
131{
132    type Output = Fut::Output;
133
134    fn run_test(self, args: TestArgs) -> Self::Output {
135        run_test(args, self)
136    }
137}
138
139impl<Fut> TestFn for fn() -> Fut
140where
141    Fut: Future,
142{
143    type Output = Fut::Output;
144
145    fn run_test(self, args: TestArgs) -> Self::Output {
146        assert!(
147            args.fixtures.is_empty(),
148            "fixtures cannot be applied for a bare function"
149        );
150        crate::rt::test_block_on(self())
151    }
152}
153
154impl TestArgs {
155    pub fn new(test_path: &'static str) -> Self {
156        TestArgs {
157            test_path,
158            migrator: None,
159            fixtures: &[],
160        }
161    }
162
163    pub fn migrator(&mut self, migrator: &'static Migrator) {
164        self.migrator = Some(migrator);
165    }
166
167    pub fn fixtures(&mut self, fixtures: &'static [TestFixture]) {
168        self.fixtures = fixtures;
169    }
170}
171
172impl TestTermination for () {
173    fn is_success(&self) -> bool {
174        true
175    }
176}
177
178impl<T, E> TestTermination for Result<T, E> {
179    fn is_success(&self) -> bool {
180        self.is_ok()
181    }
182}
183
184fn run_test_with_pool<DB, F, Fut>(args: TestArgs, test_fn: F) -> Fut::Output
185where
186    DB: TestSupport,
187    DB::Connection: Migrate,
188    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
189    F: FnOnce(Pool<DB>) -> Fut,
190    Fut: Future,
191    Fut::Output: TestTermination,
192{
193    let test_path = args.test_path;
194    run_test::<DB, _, _>(args, |pool_opts, connect_opts| async move {
195        let pool = pool_opts
196            .connect_with(connect_opts)
197            .await
198            .expect("failed to connect test pool");
199
200        let res = test_fn(pool.clone()).await;
201
202        let close_timed_out = crate::rt::timeout(Duration::from_secs(10), pool.close())
203            .await
204            .is_err();
205
206        if close_timed_out {
207            eprintln!("test {test_path} held onto Pool after exiting");
208        }
209
210        res
211    })
212}
213
214fn run_test<DB, F, Fut>(args: TestArgs, test_fn: F) -> Fut::Output
215where
216    DB: TestSupport,
217    DB::Connection: Migrate,
218    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
219    F: FnOnce(PoolOptions<DB>, <DB::Connection as Connection>::Options) -> Fut,
220    Fut: Future,
221    Fut::Output: TestTermination,
222{
223    crate::rt::test_block_on(async move {
224        let test_context = DB::test_context(&args)
225            .await
226            .expect("failed to connect to setup test database");
227
228        setup_test_db::<DB>(&test_context.connect_opts, &args).await;
229
230        let res = test_fn(test_context.pool_opts, test_context.connect_opts).await;
231
232        if res.is_success() {
233            if let Err(e) = DB::cleanup_test(&DB::db_name(&args)).await {
234                eprintln!(
235                    "failed to delete database {:?}: {}",
236                    test_context.db_name, e
237                );
238            }
239        }
240
241        res
242    })
243}
244
245async fn setup_test_db<DB: Database>(
246    copts: &<DB::Connection as Connection>::Options,
247    args: &TestArgs,
248) where
249    DB::Connection: Migrate + Sized,
250    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
251{
252    let mut conn = copts
253        .connect()
254        .await
255        .expect("failed to connect to test database");
256
257    if let Some(migrator) = args.migrator {
258        migrator
259            .run_direct(&mut conn)
260            .await
261            .expect("failed to apply migrations");
262    }
263
264    for fixture in args.fixtures {
265        (&mut conn)
266            .execute(fixture.contents)
267            .await
268            .unwrap_or_else(|e| panic!("failed to apply test fixture {:?}: {:?}", fixture.path, e));
269    }
270
271    conn.close()
272        .await
273        .expect("failed to close setup connection");
274}
sqlx_core/testing/mod.rs

sqlx_core/testing/
mod.rs
