better

Author: juleswritescode

crates/pgt_completions/src/test_helper.rs

@@ -1,7 +1,7 @@
 use std::fmt::Display;
 
 use pgt_schema_cache::SchemaCache;
-use pgt_test_utils::test_database::get_new_test_db;
+use sqlx::{Executor, PgPool};
 
 use crate::{CompletionItem, CompletionItemKind, CompletionParams, complete};
 
@@ -35,9 +35,8 @@ impl Display for InputQuery {
 pub(crate) async fn get_test_deps(
     setup: &str,
     input: InputQuery,
+    test_db: &PgPool,
 ) -> (tree_sitter::Tree, pgt_schema_cache::SchemaCache) {
-    let test_db = get_new_test_db().await;
-
     test_db
         .execute(setup)
         .await
@@ -206,8 +205,9 @@ pub(crate) async fn assert_complete_results(
     query: &str,
     assertions: Vec<CompletionAssertion>,
     setup: &str,
+    pool: &PgPool,
 ) {
-    let (tree, cache) = get_test_deps(setup, query.into()).await;
+    let (tree, cache) = get_test_deps(setup, query.into(), pool).await;
     let params = get_test_params(&tree, &cache, query.into());
     let items = complete(params);
 
@@ -240,8 +240,8 @@ pub(crate) async fn assert_complete_results(
         });
 }
 
-pub(crate) async fn assert_no_complete_results(query: &str, setup: &str) {
-    let (tree, cache) = get_test_deps(setup, query.into()).await;
+pub(crate) async fn assert_no_complete_results(query: &str, setup: &str, pool: &PgPool) {
+    let (tree, cache) = get_test_deps(setup, query.into(), pool).await;
     let params = get_test_params(&tree, &cache, query.into());
     let items = complete(params);
 

crates/pgt_lsp/tests/server.rs

@@ -13,12 +13,13 @@ use pgt_configuration::database::PartialDatabaseConfiguration;
 use pgt_fs::MemoryFileSystem;
 use pgt_lsp::LSPServer;
 use pgt_lsp::ServerFactory;
-use pgt_test_utils::test_database::get_new_test_db;
 use pgt_workspace::DynRef;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use serde_json::Value;
 use serde_json::{from_value, to_value};
+use sqlx::Executor;
+use sqlx::PgPool;
 use std::any::type_name;
 use std::fmt::Display;
 use std::time::Duration;
@@ -344,11 +345,10 @@ async fn basic_lifecycle() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_database_connection() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_database_connection(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let setup = r#"
             create table public.users (
@@ -456,11 +456,10 @@ async fn server_shutdown() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_completions() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_completions(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let setup = r#"
             create table public.users (
@@ -557,11 +556,10 @@ async fn test_completions() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_issue_271() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_issue_271(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let setup = r#"
             create table public.users (
@@ -759,11 +757,10 @@ async fn test_issue_271() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_execute_statement() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_execute_statement(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let database = test_db
         .connect_options()
@@ -898,11 +895,10 @@ async fn test_execute_statement() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_issue_281() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_issue_281(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let setup = r#"
             create table public.users (
@@ -982,11 +978,10 @@ async fn test_issue_281() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test]
-async fn test_issue_303() -> Result<()> {
+#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+async fn test_issue_303(test_db: PgPool) -> Result<()> {
     let factory = ServerFactory::default();
     let mut fs = MemoryFileSystem::default();
-    let test_db = get_new_test_db().await;
 
     let setup = r#"
             create table public.users (

crates/pgt_test_utils/src/lib.rs

@@ -1 +1,3 @@
 pub mod test_database;
+
+pub static MIGRATIONS: sqlx::migrate::Migrator = sqlx::migrate!("./testdb_migrations");

crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql

@@ -0,0 +1,23 @@
+do $$
+begin
+if not exists (
+  select from pg_catalog.pg_roles
+  where rolname = 'admin'
+) then
+  create role admin superuser createdb login bypassrls;
+end if;
+
+if not exists (
+  select from pg_catalog.pg_roles
+  where rolname = 'test_login'
+) then
+  create role test_login login;
+end if;
+
+if not exists (
+  select from pg_catalog.pg_roles
+  where rolname = 'test_nologin'
+) then
+  create role test_nologin;
+end if;
+end;