ok

Author: juleswritescode

crates/pgt_completions/src/providers/columns.rs

@@ -65,13 +65,15 @@ fn get_completion_text(ctx: &CompletionContext, func: &Function) -> CompletionTe
 
 #[cfg(test)]
 mod tests {
+    use sqlx::PgPool;
+
     use crate::{
         CompletionItem, CompletionItemKind, complete,
         test_helper::{CURSOR_POS, get_test_deps, get_test_params},
     };
 
-    #[tokio::test]
-    async fn completes_fn() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn completes_fn(pool: PgPool) {
         let setup = r#"
           create or replace function cool()
           returns trigger
@@ -86,7 +88,7 @@ mod tests {
 
         let query = format!("select coo{}", CURSOR_POS);
 
-        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await;
         let params = get_test_params(&tree, &cache, query.as_str().into());
         let results = complete(params);
 
@@ -98,8 +100,8 @@ mod tests {
         assert_eq!(label, "cool");
     }
 
-    #[tokio::test]
-    async fn prefers_fn_if_invocation() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn prefers_fn_if_invocation(pool: PgPool) {
         let setup = r#"
           create table coos (
             id serial primary key,
@@ -119,7 +121,7 @@ mod tests {
 
         let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
 
-        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await;
         let params = get_test_params(&tree, &cache, query.as_str().into());
         let results = complete(params);
 
@@ -132,8 +134,8 @@ mod tests {
         assert_eq!(kind, CompletionItemKind::Function);
     }
 
-    #[tokio::test]
-    async fn prefers_fn_in_select_clause() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn prefers_fn_in_select_clause(pool: PgPool) {
         let setup = r#"
           create table coos (
             id serial primary key,
@@ -153,7 +155,7 @@ mod tests {
 
         let query = format!(r#"select coo{}"#, CURSOR_POS);
 
-        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await;
         let params = get_test_params(&tree, &cache, query.as_str().into());
         let results = complete(params);
 
@@ -166,8 +168,8 @@ mod tests {
         assert_eq!(kind, CompletionItemKind::Function);
     }
 
-    #[tokio::test]
-    async fn prefers_function_in_from_clause_if_invocation() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn prefers_function_in_from_clause_if_invocation(pool: PgPool) {
         let setup = r#"
           create table coos (
             id serial primary key,
@@ -187,7 +189,7 @@ mod tests {
 
         let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
 
-        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let (tree, cache) = get_test_deps(Some(setup), query.as_str().into(), &pool).await;
         let params = get_test_params(&tree, &cache, query.as_str().into());
         let results = complete(params);
 

crates/pgt_completions/src/providers/functions.rs

@@ -59,10 +59,12 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi
 
 #[cfg(test)]
 mod tests {
+    use sqlx::{Executor, PgPool};
+
     use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results};
 
-    #[tokio::test]
-    async fn completes_within_quotation_marks() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn completes_within_quotation_marks(pool: PgPool) {
         let setup = r#"
             create schema private;
 
@@ -84,13 +86,16 @@ mod tests {
                 with check (true);
         "#;
 
+        pool.execute(setup).await.unwrap();
+
         assert_complete_results(
             format!("alter policy \"{}\" on private.users;", CURSOR_POS).as_str(),
             vec![
                 CompletionAssertion::Label("read for public users disallowed".into()),
                 CompletionAssertion::Label("write for public users allowed".into()),
             ],
-            setup,
+            None,
+            &pool,
         )
         .await;
 
@@ -99,7 +104,8 @@ mod tests {
             vec![CompletionAssertion::Label(
                 "write for public users allowed".into(),
             )],
-            setup,
+            None,
+            &pool,
         )
         .await;
     }

crates/pgt_completions/src/providers/policies.rs

@@ -27,13 +27,15 @@ pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut Completion
 #[cfg(test)]
 mod tests {
 
+    use sqlx::PgPool;
+
     use crate::{
         CompletionItemKind,
         test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results},
     };
 
-    #[tokio::test]
-    async fn autocompletes_schemas() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn autocompletes_schemas(pool: PgPool) {
         let setup = r#"
             create schema private;
             create schema auth;
@@ -75,13 +77,14 @@ mod tests {
                     CompletionItemKind::Schema,
                 ),
             ],
-            setup,
+            Some(setup),
+            &pool,
         )
         .await;
     }
 
-    #[tokio::test]
-    async fn suggests_tables_and_schemas_with_matching_keys() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn suggests_tables_and_schemas_with_matching_keys(pool: PgPool) {
         let setup = r#"
             create schema ultimate;
 
@@ -99,7 +102,8 @@ mod tests {
                 CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
                 CompletionAssertion::LabelAndKind("ultimate".into(), CompletionItemKind::Schema),
             ],
-            setup,
+            Some(setup),
+            &pool,
         )
         .await;
     }