sqlx_core/testing/
mod.rs

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