Update a few tests for speculative assignment

dask · Nov 24, 2020 · 2e1b5e6 · 2e1b5e6
1 parent b3b19c0
commit 2e1b5e6
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/distributed/diagnostics/progress.py b/distributed/diagnostics/progress.py
@@ -311,7 +311,14 @@ def __init__(self, scheduler):
     def create(self, key, k):
         with log_errors():
             ts = self.scheduler.tasks[key]
-            g = {"memory": 0, "erred": 0, "waiting": 0, "released": 0, "processing": 0}
+            g = {
+                "memory": 0,
+                "erred": 0,
+                "waiting": 0,
+                "released": 0,
+                "processing": 0,
+                "speculative": 0,
+            }
             self.keys[k] = set()
             self.groups[k] = g
             self.nbytes[k] = 0

diff --git a/distributed/http/scheduler/tests/test_scheduler_http.py b/distributed/http/scheduler/tests/test_scheduler_http.py
@@ -126,7 +126,15 @@ async def fetch_metrics():
         ]
         return active_metrics, forgotten_tasks
 
-    expected = {"memory", "released", "processing", "waiting", "no-worker", "erred"}
+    expected = {
+        "memory",
+        "released",
+        "processing",
+        "waiting",
+        "no-worker",
+        "erred",
+        "speculative",
+    }
 
     # Ensure that we get full zero metrics for all states even though the
     # scheduler did nothing, yet