change env var, add dotenv, load into test_db

Author: juleswritescode

.env

@@ -1 +1 @@
-DATABASE_URL=postgresql://postgres:[email protected]:5432/postgres
+DB_CONNECTION_STRING=postgresql://postgres:[email protected]:5432/postgres

Cargo.lock

@@ -63,21 +63,18 @@ pub fn complete<'a>(params: &'a CompletionParams<'a>) -> CompletionResult<'a> {
 
 #[cfg(test)]
 mod tests {
-    use async_std::task::block_on;
     use pg_schema_cache::SchemaCache;
     use pg_test_utils::test_database::*;
 
-    use sqlx::PgPool;
+    use sqlx::Executor;
 
     use crate::{complete, CompletionParams};
 
-    #[test]
-    fn test_complete() {
-        let input = "select id from c;";
-
-        let conn_string = std::env::var("DB_CONNECTION_STRING").unwrap();
+    #[tokio::test]
+    async fn test_complete() {
+        let pool = get_new_test_db().await;
 
-        let pool = block_on(PgPool::connect(conn_string.as_str())).unwrap();
+        let input = "select id from c;";
 
         let mut parser = tree_sitter::Parser::new();
         parser
@@ -86,7 +83,7 @@ mod tests {
 
         let tree = parser.parse(input, None).unwrap();
 
-        let schema_cache = block_on(SchemaCache::load(&pool));
+        let schema_cache = SchemaCache::load(&pool).await;
 
         let p = CompletionParams {
             position: 15.into(),
@@ -100,21 +97,19 @@ mod tests {
         assert!(result.items.len() > 0);
     }
 
-    #[test]
-    fn test_complete_two() {
-        let input = "select id, name, test1231234123, unknown from co;";
-
-        let conn_string = std::env::var("DB_CONNECTION_STRING").unwrap();
+    #[tokio::test]
+    async fn test_complete_two() {
+        let pool = get_new_test_db().await;
 
-        let pool = block_on(PgPool::connect(conn_string.as_str())).unwrap();
+        let input = "select id, name, test1231234123, unknown from co;";
 
         let mut parser = tree_sitter::Parser::new();
         parser
             .set_language(tree_sitter_sql::language())
             .expect("Error loading sql language");
 
         let tree = parser.parse(input, None).unwrap();
-        let schema_cache = block_on(SchemaCache::load(&pool));
+        let schema_cache = SchemaCache::load(&pool).await;
 
         let p = CompletionParams {
             position: 47.into(),
@@ -128,8 +123,10 @@ mod tests {
         assert!(result.items.len() > 0);
     }
 
-    #[test]
-    fn test_complete_with_db() {
+    #[tokio::test]
+    async fn test_complete_with_db() {
+        let test_db = get_new_test_db().await;
+
         let setup = r#"
             create table users (
                 id serial primary key,
@@ -138,23 +135,23 @@ mod tests {
             );
         "#;
 
-        let input = "select * from u";
-
-        let conn_string = std::env::var("DB_CONNECTION_STRING").unwrap();
-        let password = std::env::var("DB_PASSWORD").unwrap_or("postgres".into());
+        test_db
+            .execute(setup)
+            .await
+            .expect("Failed to execute setup query");
 
-        let test_db = block_on(get_new_test_db(conn_string, password));
+        let input = "select * from u";
 
         let mut parser = tree_sitter::Parser::new();
         parser
             .set_language(tree_sitter_sql::language())
             .expect("Error loading sql language");
 
         let tree = parser.parse(input, None).unwrap();
-        let schema_cache = block_on(SchemaCache::load(&pool));
+        let schema_cache = SchemaCache::load(&test_db).await;
 
         let p = CompletionParams {
-            position: 47.into(),
+            position: ((input.len() - 1) as u32).into(),
             schema: &schema_cache,
             text: input,
             tree: Some(&tree),

crates/pg_completions/src/lib.rs

@@ -8,3 +8,4 @@ anyhow = "1.0.81"
 uuid = { version = "1.11.0", features = ["v4"] }
 
 sqlx.workspace = true
+dotenv = "0.15.0"

crates/pg_test_utils/Cargo.toml

@@ -3,7 +3,15 @@ use uuid::Uuid;
 
 // TODO: Work with proper config objects instead of a connection_string.
 // With the current implementation, we can't parse the password from the connection string.
-pub async fn get_new_test_db(connection_string: String, database_password: String) -> PgPool {
+pub async fn get_new_test_db() -> PgPool {
+    dotenv::dotenv()
+        .ok()
+        .expect("Unable to load .env file for tests");
+
+    let connection_string =
+        std::env::var("DB_CONNECTION_STRING").expect("DB_CONNECTION_STRING not set");
+    let password = std::env::var("DB_PASSWORD").unwrap_or("postgres".into());
+
     let options_from_conn_str: PgConnectOptions = connection_string
         .parse()
         .expect("Invalid Connection String");
@@ -18,7 +26,7 @@ pub async fn get_new_test_db(connection_string: String, database_password: Strin
         .host(host)
         .port(options_from_conn_str.get_port())
         .username(options_from_conn_str.get_username())
-        .password(&database_password);
+        .password(&password);
 
     let postgres = sqlx::PgPool::connect_with(options_without_db_name.clone())
         .await