1 files changed, 102 insertions, 49 deletions

diff --git a/test_util/src/lsp.rs b/test_util/src/lsp.rs
index a7061543f..831df28e9 100644
--- a/test_util/src/lsp.rs
+++ b/test_util/src/lsp.rs
@@ -87,6 +87,12 @@ impl<'a> From<&'a [u8]> for LspMessage {
   }
 }
 
+#[derive(Debug, Deserialize)]
+struct DiagnosticBatchNotificationParams {
+  batch_index: usize,
+  messages_len: usize,
+}
+
 fn read_message<R>(reader: &mut R) -> Result<Option<Vec<u8>>>
 where
   R: io::Read + io::BufRead,
@@ -174,6 +180,25 @@ impl LspStdoutReader {
       cvar.wait(&mut msg_queue);
     }
   }
+
+  pub fn read_latest_message<R>(
+    &mut self,
+    mut get_match: impl FnMut(&LspMessage) -> Option<R>,
+  ) -> R {
+    let (msg_queue, cvar) = &*self.pending_messages;
+    let mut msg_queue = msg_queue.lock();
+    loop {
+      for i in (0..msg_queue.len()).rev() {
+        let msg = &msg_queue[i];
+        if let Some(result) = get_match(msg) {
+          let msg = msg_queue.remove(i);
+          self.read_messages.push(msg);
+          return result;
+        }
+      }
+      cvar.wait(&mut msg_queue);
+    }
+  }
 }
 
 pub struct InitializeParamsBuilder {
@@ -445,6 +470,7 @@ pub struct LspClientBuilder {
   print_stderr: bool,
   deno_exe: PathBuf,
   context: Option<TestContext>,
+  use_diagnostic_sync: bool,
 }
 
 impl LspClientBuilder {
@@ -454,6 +480,7 @@ impl LspClientBuilder {
       print_stderr: false,
       deno_exe: deno_exe_path(),
       context: None,
+      use_diagnostic_sync: true,
     }
   }
 
@@ -470,6 +497,13 @@ impl LspClientBuilder {
     self
   }
 
+  /// Whether to use the synchronization messages to better sync diagnostics
+  /// between the test client and server.
+  pub fn use_diagnostic_sync(&mut self, value: bool) -> &mut Self {
+    self.use_diagnostic_sync = value;
+    self
+  }
+
   pub fn set_test_context(&mut self, test_context: &TestContext) -> &mut Self {
     self.context = Some(test_context.clone());
     self
@@ -485,6 +519,11 @@ impl LspClientBuilder {
     command
       .env("DENO_DIR", deno_dir.path())
       .env("NPM_CONFIG_REGISTRY", npm_registry_url())
+      // turn on diagnostic synchronization communication
+      .env(
+        "DENO_DONT_USE_INTERNAL_LSP_DIAGNOSTIC_SYNC_FLAG",
+        if self.use_diagnostic_sync { "1" } else { "" },
+      )
       .arg("lsp")
       .stdin(Stdio::piped())
       .stdout(Stdio::piped());
@@ -510,7 +549,6 @@ impl LspClientBuilder {
         .unwrap_or_else(|| TestContextBuilder::new().build()),
       writer,
       deno_dir,
-      diagnosable_open_file_count: 0,
     })
   }
 }
@@ -523,7 +561,6 @@ pub struct LspClient {
   writer: io::BufWriter<ChildStdin>,
   deno_dir: TempDir,
   context: TestContext,
-  diagnosable_open_file_count: usize,
 }
 
 impl Drop for LspClient {
@@ -609,20 +646,6 @@ impl LspClient {
   }
 
   pub fn did_open_raw(&mut self, params: Value) {
-    let text_doc = params
-      .as_object()
-      .unwrap()
-      .get("textDocument")
-      .unwrap()
-      .as_object()
-      .unwrap();
-    if matches!(
-      text_doc.get("languageId").unwrap().as_str().unwrap(),
-      "typescript" | "javascript"
-    ) {
-      self.diagnosable_open_file_count += 1;
-    }
-
     self.write_notification("textDocument/didOpen", params);
   }
 
@@ -632,11 +655,46 @@ impl LspClient {
     self.write_response(id, result);
   }
 
+  fn get_latest_diagnostic_batch_index(&mut self) -> usize {
+    let result = self
+      .write_request("deno/internalLatestDiagnosticBatchIndex", json!(null));
+    result.as_u64().unwrap() as usize
+  }
+
+  /// Reads the latest diagnostics. It's assumed that
   pub fn read_diagnostics(&mut self) -> CollectedDiagnostics {
-    let mut all_diagnostics = Vec::new();
-    for _ in 0..self.diagnosable_open_file_count {
-      all_diagnostics.extend(read_diagnostics(self).0);
+    // ask the server what the latest diagnostic batch index is
+    let latest_diagnostic_batch_index =
+      self.get_latest_diagnostic_batch_index();
+
+    // now wait for three (deno, lint, and typescript diagnostics) batch
+    // notification messages for that index
+    let mut read = 0;
+    let mut total_messages_len = 0;
+    while read < 3 {
+      let (method, response) =
+        self.read_notification::<DiagnosticBatchNotificationParams>();
+      assert_eq!(method, "deno/internalTestDiagnosticBatch");
+      let response = response.unwrap();
+      if response.batch_index == latest_diagnostic_batch_index {
+        read += 1;
+        total_messages_len += response.messages_len;
+      }
     }
+
+    // now read the latest diagnostic messages
+    let mut all_diagnostics = Vec::with_capacity(total_messages_len);
+    let mut seen_files = HashSet::new();
+    for _ in 0..total_messages_len {
+      let (method, response) =
+        self.read_latest_notification::<lsp::PublishDiagnosticsParams>();
+      assert_eq!(method, "textDocument/publishDiagnostics");
+      let response = response.unwrap();
+      if seen_files.insert(response.uri.to_string()) {
+        all_diagnostics.push(response);
+      }
+    }
+
     CollectedDiagnostics(all_diagnostics)
   }
 
@@ -668,6 +726,19 @@ impl LspClient {
     })
   }
 
+  pub fn read_latest_notification<R>(&mut self) -> (String, Option<R>)
+  where
+    R: de::DeserializeOwned,
+  {
+    self.reader.read_latest_message(|msg| match msg {
+      LspMessage::Notification(method, maybe_params) => {
+        let params = serde_json::from_value(maybe_params.clone()?).ok()?;
+        Some((method.to_string(), params))
+      }
+      _ => None,
+    })
+  }
+
   pub fn read_notification_with_method<R>(
     &mut self,
     expected_method: &str,
@@ -819,35 +890,29 @@ impl LspClient {
 }
 
 #[derive(Debug, Clone)]
-pub struct CollectedDiagnostics(pub Vec<lsp::PublishDiagnosticsParams>);
+pub struct CollectedDiagnostics(Vec<lsp::PublishDiagnosticsParams>);
 
 impl CollectedDiagnostics {
   /// Gets the diagnostics that the editor will see after all the publishes.
-  pub fn viewed(&self) -> Vec<lsp::Diagnostic> {
+  pub fn all(&self) -> Vec<lsp::Diagnostic> {
     self
-      .viewed_messages()
+      .all_messages()
       .into_iter()
       .flat_map(|m| m.diagnostics)
       .collect()
   }
 
   /// Gets the messages that the editor will see after all the publishes.
-  pub fn viewed_messages(&self) -> Vec<lsp::PublishDiagnosticsParams> {
-    // go over the publishes in reverse order in order to get
-    // the final messages that will be shown in the editor
-    let mut messages = Vec::new();
-    let mut had_specifier = HashSet::new();
-    for message in self.0.iter().rev() {
-      if had_specifier.insert(message.uri.clone()) {
-        messages.insert(0, message.clone());
-      }
-    }
-    messages
+  pub fn all_messages(&self) -> Vec<lsp::PublishDiagnosticsParams> {
+    self.0.clone()
   }
 
-  pub fn with_source(&self, source: &str) -> lsp::PublishDiagnosticsParams {
+  pub fn messages_with_source(
+    &self,
+    source: &str,
+  ) -> lsp::PublishDiagnosticsParams {
     self
-      .viewed_messages()
+      .all_messages()
       .iter()
       .find(|p| {
         p.diagnostics
@@ -858,14 +923,14 @@ impl CollectedDiagnostics {
       .unwrap()
   }
 
-  pub fn with_file_and_source(
+  pub fn messages_with_file_and_source(
     &self,
     specifier: &str,
     source: &str,
   ) -> lsp::PublishDiagnosticsParams {
     let specifier = Url::parse(specifier).unwrap();
     self
-      .viewed_messages()
+      .all_messages()
       .iter()
       .find(|p| {
         p.uri == specifier
@@ -879,18 +944,6 @@ impl CollectedDiagnostics {
   }
 }
 
-fn read_diagnostics(client: &mut LspClient) -> CollectedDiagnostics {
-  // diagnostics come in batches of three unless they're cancelled
-  let mut diagnostics = vec![];
-  for _ in 0..3 {
-    let (method, response) =
-      client.read_notification::<lsp::PublishDiagnosticsParams>();
-    assert_eq!(method, "textDocument/publishDiagnostics");
-    diagnostics.push(response.unwrap());
-  }
-  CollectedDiagnostics(diagnostics)
-}
-
 #[cfg(test)]
 mod tests {
   use super::*;