def test_get_task_duration():
with dask.config.set(
{"distributed.scheduler.default-task-durations": {"prefix_1": 100}}
):
s = Scheduler(port=0)
assert "prefix_1" in s.task_duration
assert s.task_duration["prefix_1"] == 100
ts_pref1 = TaskState("prefix_1-abcdefab", None)
assert s.get_task_duration(ts_pref1) == 100
# make sure get_task_duration adds TaskStates to unknown dict
assert len(s.unknown_durations) == 0
ts_pref2 = TaskState("prefix_2-abcdefab", None)
assert s.get_task_duration(ts_pref2) == 0.5 # default
assert len(s.unknown_durations) == 1
assert len(s.unknown_durations["prefix_2"]) == 1
ts_pref2_2 = TaskState("prefix_2-accdefab", None)
assert s.get_task_duration(ts_pref2_2) == 0.5 # default
assert len(s.unknown_durations) == 1
assert len(s.unknown_durations["prefix_2"]) == 2
