1mod buffer_tests;
2mod db_tests;
3
4use super::*;
5use gpui::executor::Background;
6use parking_lot::Mutex;
7use sea_orm::ConnectionTrait;
8use sqlx::migrate::MigrateDatabase;
9use std::sync::Arc;
10
11pub struct TestDb {
12 pub db: Option<Arc<Database>>,
13 pub connection: Option<sqlx::AnyConnection>,
14}
15
16impl TestDb {
17 pub fn sqlite(background: Arc<Background>) -> Self {
18 let url = format!("sqlite::memory:");
19 let runtime = tokio::runtime::Builder::new_current_thread()
20 .enable_io()
21 .enable_time()
22 .build()
23 .unwrap();
24
25 let mut db = runtime.block_on(async {
26 let mut options = ConnectOptions::new(url);
27 options.max_connections(5);
28 let db = Database::new(options, Executor::Deterministic(background))
29 .await
30 .unwrap();
31 let sql = include_str!(concat!(
32 env!("CARGO_MANIFEST_DIR"),
33 "/migrations.sqlite/20221109000000_test_schema.sql"
34 ));
35 db.pool
36 .execute(sea_orm::Statement::from_string(
37 db.pool.get_database_backend(),
38 sql.into(),
39 ))
40 .await
41 .unwrap();
42 db
43 });
44
45 db.runtime = Some(runtime);
46
47 Self {
48 db: Some(Arc::new(db)),
49 connection: None,
50 }
51 }
52
53 pub fn postgres(background: Arc<Background>) -> Self {
54 static LOCK: Mutex<()> = Mutex::new(());
55
56 let _guard = LOCK.lock();
57 let mut rng = StdRng::from_entropy();
58 let url = format!(
59 "postgres://postgres@localhost/zed-test-{}",
60 rng.gen::<u128>()
61 );
62 let runtime = tokio::runtime::Builder::new_current_thread()
63 .enable_io()
64 .enable_time()
65 .build()
66 .unwrap();
67
68 let mut db = runtime.block_on(async {
69 sqlx::Postgres::create_database(&url)
70 .await
71 .expect("failed to create test db");
72 let mut options = ConnectOptions::new(url);
73 options
74 .max_connections(5)
75 .idle_timeout(Duration::from_secs(0));
76 let db = Database::new(options, Executor::Deterministic(background))
77 .await
78 .unwrap();
79 let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
80 db.migrate(Path::new(migrations_path), false).await.unwrap();
81 db
82 });
83
84 db.runtime = Some(runtime);
85
86 Self {
87 db: Some(Arc::new(db)),
88 connection: None,
89 }
90 }
91
92 pub fn db(&self) -> &Arc<Database> {
93 self.db.as_ref().unwrap()
94 }
95}
96
97#[macro_export]
98macro_rules! test_both_dbs {
99 ($test_name:ident, $postgres_test_name:ident, $sqlite_test_name:ident) => {
100 #[gpui::test]
101 async fn $postgres_test_name() {
102 let test_db = crate::db::TestDb::postgres(
103 gpui::executor::Deterministic::new(0).build_background(),
104 );
105 $test_name(test_db.db()).await;
106 }
107
108 #[gpui::test]
109 async fn $sqlite_test_name() {
110 let test_db =
111 crate::db::TestDb::sqlite(gpui::executor::Deterministic::new(0).build_background());
112 $test_name(test_db.db()).await;
113 }
114 };
115}
116
117impl Drop for TestDb {
118 fn drop(&mut self) {
119 let db = self.db.take().unwrap();
120 if let sea_orm::DatabaseBackend::Postgres = db.pool.get_database_backend() {
121 db.runtime.as_ref().unwrap().block_on(async {
122 use util::ResultExt;
123 let query = "
124 SELECT pg_terminate_backend(pg_stat_activity.pid)
125 FROM pg_stat_activity
126 WHERE
127 pg_stat_activity.datname = current_database() AND
128 pid <> pg_backend_pid();
129 ";
130 db.pool
131 .execute(sea_orm::Statement::from_string(
132 db.pool.get_database_backend(),
133 query.into(),
134 ))
135 .await
136 .log_err();
137 sqlx::Postgres::drop_database(db.options.get_url())
138 .await
139 .log_err();
140 })
141 }
142 }
143}