tries to fix tests

Author: dearfl

pyrefly/lib/test/lsp/lsp_interaction/diagnostic.rs

@@ -14,6 +14,7 @@ use pyrefly_config::environment::environment::PythonEnvironment;
 
 use crate::test::lsp::lsp_interaction::object_model::InitializeSettings;
 use crate::test::lsp::lsp_interaction::object_model::LspInteraction;
+use crate::test::lsp::lsp_interaction::object_model::ValidationResult;
 use crate::test::lsp::lsp_interaction::util::get_test_files_root;
 
 #[test]
@@ -663,65 +664,85 @@ fn test_version_support_publish_diagnostics() {
         ..Default::default()
     });
 
+    let gen_validator = |expected_version: i64| {
+        let actual_uri = uri.as_str();
+        move |msg: &Message| {
+            let Message::Notification(Notification { method, params }) = msg else {
+                return ValidationResult::Skip;
+            };
+            let Some(uri_val) = params.get("uri") else {
+                return ValidationResult::Skip;
+            };
+            let Some(expected_uri) = uri_val.as_str() else {
+                return ValidationResult::Skip;
+            };
+            if expected_uri == actual_uri && method == "textDocument/publishDiagnostics" {
+                if let Some(actual_version) = params.get("version") {
+                    if let Some(actual_version) = actual_version.as_i64() {
+                        assert!(
+                            actual_version <= expected_version,
+                            "expected version: {}, actual version: {}",
+                            expected_version,
+                            actual_version
+                        );
+                        return match actual_version.cmp(&expected_version) {
+                            std::cmp::Ordering::Less => ValidationResult::Skip,
+                            std::cmp::Ordering::Equal => ValidationResult::Pass,
+                            std::cmp::Ordering::Greater => ValidationResult::Fail,
+                        };
+                    }
+                }
+            }
+            ValidationResult::Skip
+        }
+    };
+
     interaction.server.did_open("text_document.py");
-    interaction.server.diagnostic("text_document.py");
 
-    interaction
-        .client
-        .expect_message(lsp_server::Message::Notification(
-            lsp_server::Notification {
-                method: "textDocument/publishDiagnostics".to_owned(),
-                params: serde_json::json! {{
-                    "uri": uri,
-                    "diagnostics": [],
-                    "version": 1
-                }},
-            },
-        ));
+    let version = 1;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
 
-    interaction.server.did_change("text_document.py", "# test");
-    interaction.server.diagnostic("text_document.py");
+    interaction.server.did_change("text_document.py", "a = b");
 
-    // I don't understand why this version is still 1
-    interaction
-        .client
-        .expect_message(lsp_server::Message::Notification(
-            lsp_server::Notification {
-                method: "textDocument/publishDiagnostics".to_owned(),
-                params: serde_json::json! {{
-                    "uri": uri,
-                    "diagnostics": [],
-                    "version": 2
-                }},
-            },
-        ));
+    let version = 2;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
 
     interaction
         .server
         .send_message(Message::Notification(Notification {
             method: "textDocument/didClose".to_owned(),
             params: serde_json::json!({
                 "textDocument": {
-                    "uri": uri.to_string(),
+                    "uri": uri.as_str(),
                     "languageId": "python",
                     "version": 3
                 },
             }),
         }));
-    interaction.server.diagnostic("text_document.py");
 
-    interaction
-        .client
-        .expect_message(lsp_server::Message::Notification(
-            lsp_server::Notification {
-                method: "textDocument/publishDiagnostics".to_owned(),
-                params: serde_json::json! {{
-                    "uri": uri,
-                    "diagnostics": [],
-                    "version": 3
-                }},
-            },
-        ));
+    let version = 3;
+    interaction.client.expect_message_helper(
+        gen_validator(version),
+        &format!(
+            "publishDiagnostics notification with version {} for file: {}",
+            version,
+            uri.as_str()
+        ),
+    );
 
     interaction.shutdown();
 }