ok

Author: juleswritescode

crates/pgt_completions/src/providers/roles.rs

@@ -27,48 +27,56 @@ pub fn complete_roles<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionB
 
 #[cfg(test)]
 mod tests {
+    use sqlx::PgPool;
+
     use crate::test_helper::{CURSOR_POS, CompletionAssertion, assert_complete_results};
 
     const SETUP: &'static str = r#"
-            do $$
-            begin
-                if not exists (
-                    select from pg_catalog.pg_roles
-                    where rolname = 'test'
-                ) then
-                    create role test;
-                end if;
-            end $$;
-
             create table users (
               id serial primary key,
               email varchar,
               address text
             );
         "#;
 
-    #[tokio::test]
-    async fn works_in_drop_role() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn works_in_drop_role(pool: PgPool) {
         assert_complete_results(
             format!("drop role {}", CURSOR_POS).as_str(),
-            vec![CompletionAssertion::LabelAndKind(
-                "test".into(),
-                crate::CompletionItemKind::Role,
-            )],
-            SETUP,
+            vec![
+                CompletionAssertion::LabelAndKind("admin".into(), crate::CompletionItemKind::Role),
+                CompletionAssertion::LabelAndKind(
+                    "test_login".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+                CompletionAssertion::LabelAndKind(
+                    "test_nologin".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+            ],
+            Some(SETUP),
+            &pool,
         )
         .await;
     }
 
-    #[tokio::test]
-    async fn works_in_alter_role() {
+    #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+    async fn works_in_alter_role(pool: PgPool) {
         assert_complete_results(
             format!("alter role {}", CURSOR_POS).as_str(),
-            vec![CompletionAssertion::LabelAndKind(
-                "test".into(),
-                crate::CompletionItemKind::Role,
-            )],
-            SETUP,
+            vec![
+                CompletionAssertion::LabelAndKind("admin".into(), crate::CompletionItemKind::Role),
+                CompletionAssertion::LabelAndKind(
+                    "test_login".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+                CompletionAssertion::LabelAndKind(
+                    "test_nologin".into(),
+                    crate::CompletionItemKind::Role,
+                ),
+            ],
+            Some(SETUP),
+            &pool,
         )
         .await;
     }

crates/pgt_completions/src/relevance/scoring.rs

@@ -208,7 +208,7 @@ impl CompletionScore<'_> {
         }
     }
 
-    fn get_schema_name(&self) -> &str {
+    fn get_schema_name(&self) -> Option<&str> {
         match self.data {
             CompletionRelevanceData::Function(f) => Some(f.schema.as_str()),
             CompletionRelevanceData::Table(t) => Some(t.schema.as_str()),
@@ -283,7 +283,7 @@ impl CompletionScore<'_> {
             return;
         }
 
-        let schema = match self.get_schema_name() {
+        let schema_name = match self.get_schema_name() {
             Some(s) => s.to_string(),
             None => return,
         };

crates/pgt_schema_cache/src/roles.rs

@@ -33,7 +33,7 @@ mod tests {
 
         let roles = &cache.roles;
 
-        let super_role = roles.iter().find(|r| r.name == "test_super").unwrap();
+        let super_role = roles.iter().find(|r| r.name == "admin").unwrap();
         assert!(super_role.is_super_user);
         assert!(super_role.can_create_db);
         assert!(super_role.can_login);

crates/pgt_test_utils/testdb_migrations/0001-setup-roles.sql renamed to crates/pgt_test_utils/testdb_migrations/0001_setup-roles.sql

@@ -1,5 +1,6 @@
 do $$
 begin
+
 if not exists (
   select from pg_catalog.pg_roles
   where rolname = 'admin'
@@ -20,4 +21,6 @@ if not exists (
 ) then
   create role test_nologin;
 end if;
-end;
+
+end
+$$;