From d7934d8e2089a314a81936da5776511abf50a96e Mon Sep 17 00:00:00 2001
From: crusaderky <crusaderky@gmail.com>
Date: Mon, 16 May 2022 18:00:01 +0100
Subject: [PATCH] Don't manually tamper with the state

---
 distributed/tests/test_worker_state_machine.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
index 51ca129b7da..6ec1ea4dc04 100644
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -427,11 +427,12 @@ async def test_self_denounce_missing_data(c, s, a):
     x = c.submit(inc, 1, key="x")
     await x
 
-    # Manually wipe x from the worker
-    # (Using an endpoint like remove-replicas would inform the scheduler).
-    del a.data["x"]
-    del a.tasks["x"]
-    a.validate_state()
+    # Wipe x from the worker and lose the message that would inform the scheduler
+    a.handle_remove_replicas(keys=["x"], stimulus_id="test")
+    assert a.batched_stream.buffer == [
+        {"key": "x", "stimulus_id": "test", "op": "release-worker-data"}
+    ]
+    a.batched_stream.buffer.clear()
 
     y = c.submit(inc, x, key="y")
     # The scheduler tries computing y, but a responds that x is not available.