ensure_default_client updates state on client

fjetter · fjetter · commit 7fd04cfb782a · 2022-04-01T13:17:42.000+02:00
diff --git a/distributed/client.py b/distributed/client.py
@@ -1516,7 +1516,8 @@ async def _close(self, fast=False):
             return
 
         self.status = "closing"
-        _del_global_client(self)
+        if self._set_as_default:
+            _del_global_client(self)
 
         for preload in self.preloads:
             await preload.teardown()
@@ -5201,6 +5202,7 @@ def ensure_default_client(client: Client) -> None:
     client : Client
         The client
     """
+    client._set_as_default = True
     _set_global_client(client)
 
 
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
@@ -3461,7 +3461,9 @@ async def test_default_global_client_multi_clients(s):
 async def test_ensure_default_client(c, s, a, b):
     assert c is default_client()
 
-    async with Client(s.address, set_as_default=False, asynchronous=True) as c2:
+    async with Client(
+        s.address, set_as_default=False, asynchronous=True, name="c2"
+    ) as c2:
         assert c is default_client()
         assert c2 is not default_client()
         ensure_default_client(c2)
