chore: more tests (#218)

Author: psteinroe

Cargo.lock

@@ -23,6 +23,7 @@ crossbeam                = "0.8.4"
 enumflags2               = "0.7.10"
 ignore                   = "0.4.23"
 indexmap                 = { version = "2.6.0", features = ["serde"] }
+insta                    = "1.31.0"
 line_index               = { path = "./lib/line_index", version = "0.0.0" }
 pg_query                 = "6.0.0"
 proc-macro2              = "1.0.66"

Cargo.toml

@@ -10,7 +10,7 @@ pub(crate) fn init(mut session: CliSession) -> Result<(), CliDiagnostic> {
     let file_created = ConfigName::pglt_toml();
     session.app.console.log(markup! {
 "
-Welcome to the Postgres Language Server! Let's get you started...
+Welcome to the Postgres Language Tools! Let's get you started...
 
 "<Info><Emphasis>"Files created "</Emphasis></Info>"
 

crates/pglt_cli/src/commands/init.rs

@@ -21,7 +21,7 @@ cstree              = { version = "0.12.0", features = ["derive"] }
 text-size.workspace = true
 
 [dev-dependencies]
-insta = "1.31.0"
+insta.workspace = true
 
 [lib]
 doctest = false

crates/pglt_lexer/Cargo.toml

@@ -33,9 +33,11 @@ tower-lsp           = { version = "0.20.0" }
 tracing             = { workspace = true, features = ["attributes"] }
 
 [dev-dependencies]
-tokio = { workspace = true, features = ["macros"] }
-tower = { version = "0.4.13", features = ["timeout"] }
-
+pglt_test_utils = { workspace = true }
+sqlx            = { workspace = true }
+tokio           = { workspace = true, features = ["macros"] }
+toml            = { workspace = true }
+tower           = { version = "0.4.13", features = ["timeout"] }
 
 [lib]
 doctest = false

crates/pglt_lsp/Cargo.toml

@@ -2,19 +2,28 @@ use anyhow::bail;
 use anyhow::Context;
 use anyhow::Error;
 use anyhow::Result;
+use biome_deserialize::Merge;
 use futures::channel::mpsc::{channel, Sender};
 use futures::Sink;
 use futures::SinkExt;
 use futures::Stream;
 use futures::StreamExt;
+use pglt_configuration::database::PartialDatabaseConfiguration;
+use pglt_configuration::PartialConfiguration;
+use pglt_fs::MemoryFileSystem;
 use pglt_lsp::LSPServer;
 use pglt_lsp::ServerFactory;
+use pglt_test_utils::test_database::get_new_test_db;
+use pglt_workspace::DynRef;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use serde_json::{from_value, to_value};
+use sqlx::Executor;
 use std::any::type_name;
 use std::fmt::Display;
+use std::process::id;
 use std::time::Duration;
+use tokio::time::sleep;
 use tower::timeout::Timeout;
 use tower::{Service, ServiceExt};
 use tower_lsp::jsonrpc;
@@ -314,3 +323,115 @@ async fn basic_lifecycle() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_database_connection() -> 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 (
+                id serial primary key,
+                name varchar(255) not null
+            );
+        "#;
+
+    test_db
+        .execute(setup)
+        .await
+        .expect("Failed to setup test database");
+
+    let mut conf = PartialConfiguration::init();
+    conf.merge_with(PartialConfiguration {
+        db: Some(PartialDatabaseConfiguration {
+            database: Some(
+                test_db
+                    .connect_options()
+                    .get_database()
+                    .unwrap()
+                    .to_string(),
+            ),
+            ..Default::default()
+        }),
+        ..Default::default()
+    });
+    fs.insert(
+        url!("pglt.toml").to_file_path().unwrap(),
+        toml::to_string(&conf).unwrap(),
+    );
+
+    let (service, client) = factory
+        .create_with_fs(None, DynRef::Owned(Box::new(fs)))
+        .into_inner();
+
+    let (stream, sink) = client.split();
+    let mut server = Server::new(service);
+
+    let (sender, mut receiver) = channel(CHANNEL_BUFFER_SIZE);
+    let reader = tokio::spawn(client_handler(stream, sink, sender));
+
+    server.initialize().await?;
+    server.initialized().await?;
+
+    server.load_configuration().await?;
+
+    server
+        .open_document("select unknown from public.users; ")
+        .await?;
+
+    // in this test, we want to ensure a database connection is established and the schema cache is
+    // loaded. This is the case when the server sends typecheck diagnostics for the query above.
+    // so we wait for diagnostics to be sent.
+    let notification = tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            match receiver.next().await {
+                Some(ServerNotification::PublishDiagnostics(msg)) => {
+                    if msg
+                        .diagnostics
+                        .iter()
+                        .any(|d| d.message.contains("column \"unknown\" does not exist"))
+                    {
+                        return true;
+                    }
+                }
+                _ => continue,
+            }
+        }
+    })
+    .await
+    .is_ok();
+
+    assert!(notification, "expected diagnostics for unknown column");
+
+    server.shutdown().await?;
+    reader.abort();
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn server_shutdown() -> Result<()> {
+    let factory = ServerFactory::default();
+    let (service, client) = factory.create(None).into_inner();
+    let (stream, sink) = client.split();
+    let mut server = Server::new(service);
+
+    let (sender, _) = channel(CHANNEL_BUFFER_SIZE);
+    let reader = tokio::spawn(client_handler(stream, sink, sender));
+
+    server.initialize().await?;
+    server.initialized().await?;
+
+    let cancellation = factory.cancellation();
+    let cancellation = cancellation.notified();
+
+    // this is called when `pglt stop` is run by the user
+    server.pglt_shutdown().await?;
+
+    cancellation.await;
+
+    reader.abort();
+
+    Ok(())
+}

crates/pglt_lsp/tests/server.rs

@@ -12,7 +12,6 @@ version              = "0.0.0"
 
 
 [dependencies]
-insta                       = "1.31.0"
 pglt_console.workspace      = true
 pglt_diagnostics.workspace  = true
 pglt_query_ext.workspace    = true
@@ -24,6 +23,7 @@ tree-sitter.workspace       = true
 tree_sitter_sql.workspace   = true
 
 [dev-dependencies]
+insta.workspace           = true
 pglt_test_utils.workspace = true
 
 [lib]