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 fn test_context(args: &TestArgs) -> BoxFuture<'_, Result<TestContext<Self>, Error>>;
28
29 fn cleanup_test(db_name: &str) -> BoxFuture<'_, Result<(), Error>>;
30
31 fn cleanup_test_dbs() -> BoxFuture<'static, Result<Option<usize>, Error>>;
38
39 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}