From f8d646794d64577c7f75279a06491f7bc2b7d918 Mon Sep 17 00:00:00 2001
From: Kirill Bulatov <kirill@zed.dev>
Date: Thu, 2 Apr 2026 16:53:37 +0300
Subject: [PATCH] Fixup and test edit prediction mocks (#52991)

Follow-up of https://github.com/zed-industries/zed/pull/40185

Release Notes:

- N/A
---
 crates/editor/src/edit_prediction_tests.rs | 40 ++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/crates/editor/src/edit_prediction_tests.rs b/crates/editor/src/edit_prediction_tests.rs
index d1e326bc93b8052f3ae089c211e65eb3ef020fdf..987801471e5602f256ce2dd65edd57873c878027 100644
--- a/crates/editor/src/edit_prediction_tests.rs
+++ b/crates/editor/src/edit_prediction_tests.rs
@@ -1081,6 +1081,44 @@ async fn test_cancel_clears_stale_edit_prediction_in_menu(cx: &mut gpui::TestApp
     });
 }
 
+#[gpui::test]
+async fn test_discard_clears_delegate_completion(cx: &mut gpui::TestAppContext) {
+    init_test(cx, |_| {});
+    load_default_keymap(cx);
+
+    let mut cx = EditorTestContext::new(cx).await;
+    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
+    assign_editor_completion_provider(provider.clone(), &mut cx);
+    cx.set_state("let x = ˇ;");
+
+    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
+    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+
+    cx.update_editor(|editor, _window, _cx| {
+        assert!(editor.active_edit_prediction.is_some());
+    });
+
+    // Dismiss the prediction — this must call discard() on the delegate,
+    // which should clear self.completion.
+    cx.simulate_keystroke("escape");
+    cx.run_until_parked();
+
+    cx.update_editor(|editor, _window, _cx| {
+        assert!(editor.active_edit_prediction.is_none());
+    });
+
+    // update_visible_edit_prediction must NOT bring the prediction back,
+    // because discard() cleared self.completion in the delegate.
+    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+
+    cx.update_editor(|editor, _window, _cx| {
+        assert!(
+            editor.active_edit_prediction.is_none(),
+            "prediction must not resurface after discard()"
+        );
+    });
+}
+
 fn accept_completion(cx: &mut EditorTestContext) {
     cx.update_editor(|editor, window, cx| {
         editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
@@ -1350,6 +1388,7 @@ impl EditPredictionDelegate for FakeEditPredictionDelegate {
         _reason: edit_prediction_types::EditPredictionDiscardReason,
         _cx: &mut gpui::Context<Self>,
     ) {
+        self.completion.take();
     }
 
     fn suggest<'a>(
@@ -1426,6 +1465,7 @@ impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
         _reason: edit_prediction_types::EditPredictionDiscardReason,
         _cx: &mut gpui::Context<Self>,
     ) {
+        self.completion.take();
     }
 
     fn suggest<'a>(