Fix regressions

dask · Oct 20, 2022 · 0850804 · 0850804
1 parent 998de6c
commit 0850804
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 7 deletions.
diff --git a/distributed/http/worker/tests/test_worker_http.py b/distributed/http/worker/tests/test_worker_http.py
@@ -71,8 +71,14 @@ async def fetch_state_metrics():
     assert not a.state.tasks
     active_metrics = await fetch_state_metrics()
     assert active_metrics == {
-        "stored": 0.0,
+        "constrained": 0.0,
         "executing": 0.0,
+        "fetch": 0.0,
+        "flight": 0.0,
+        "long-running": 0.0,
+        "memory": 0.0,
+        "missing": 0.0,
+        "other": 0.0,
         "ready": 0.0,
         "waiting": 0.0,
     }
@@ -86,8 +92,14 @@ async def fetch_state_metrics():
 
     active_metrics = await fetch_state_metrics()
     assert active_metrics == {
-        "stored": 0.0,
+        "constrained": 0.0,
         "executing": 1.0,
+        "fetch": 0.0,
+        "flight": 0.0,
+        "long-running": 0.0,
+        "memory": 0.0,
+        "missing": 0.0,
+        "other": 0.0,
         "ready": 0.0,
         "waiting": 0.0,
     }
@@ -102,8 +114,14 @@ async def fetch_state_metrics():
 
     active_metrics = await fetch_state_metrics()
     assert active_metrics == {
-        "stored": 0.0,
+        "constrained": 0.0,
         "executing": 0.0,
+        "fetch": 0.0,
+        "flight": 0.0,
+        "long-running": 0.0,
+        "memory": 0.0,
+        "missing": 0.0,
+        "other": 0.0,
         "ready": 0.0,
         "waiting": 0.0,
     }

diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
@@ -703,10 +703,9 @@ async def test_override_data_worker(s):
     async with Worker(s.address, data=UserDict) as w:
         assert type(w.data) is UserDict
 
-    data = UserDict({"x": 1})
+    data = UserDict()
     async with Worker(s.address, data=data) as w:
         assert w.data is data
-        assert w.data == {"x": 1}
 
 
 @gen_cluster(

diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
@@ -1676,8 +1676,24 @@ def test_fetch_count(ws):
     assert len(ws.missing_dep_flight) == 1
 
 
-def test_task_counts(ws_with_running_task):
-    ws = ws_with_running_task
+def test_task_counts(ws):
+    assert ws.task_counts == {
+        "constrained": 0,
+        "executing": 0,
+        "fetch": 0,
+        "flight": 0,
+        "long-running": 0,
+        "memory": 0,
+        "missing": 0,
+        "other": 0,
+        "ready": 0,
+        "waiting": 0,
+    }
+
+
+def test_task_counts_actor_in_memory(ws):
+    ws.handle_stimulus(ComputeTaskEvent.dummy("x", actor=True, stimulus_id="s1"))
+    assert ws.actors == {"x": None}
     assert ws.task_counts == {
         "constrained": 0,
         "executing": 1,
@@ -1690,3 +1706,17 @@ def test_task_counts(ws_with_running_task):
         "ready": 0,
         "waiting": 0,
     }
+    ws.handle_stimulus(ExecuteSuccessEvent.dummy("x", value=123, stimulus_id="s2"))
+    assert ws.actors == {"x": 123}
+    assert ws.task_counts == {
+        "constrained": 0,
+        "executing": 0,
+        "fetch": 0,
+        "flight": 0,
+        "long-running": 0,
+        "memory": 1,
+        "missing": 0,
+        "other": 0,
+        "ready": 0,
+        "waiting": 0,
+    }