Fix test_internal_kv CI

tests/test_internal_kv.py

@@ -6,28 +6,29 @@
 import fed._private.compatible_utils as compatible_utils
 
 
-def test_kv_init():
-    def run(party):
-        compatible_utils.init_ray("local")
-        cluster = {
-            'alice': {'address': '127.0.0.1:11010', 'listen_addr': '0.0.0.0:11010'},
-            'bob': {'address': '127.0.0.1:11011', 'listen_addr': '0.0.0.0:11011'},
-        }
-        assert compatible_utils.kv is None
-        fed.init(cluster=cluster, party=party)
-        assert compatible_utils.kv
-        assert not compatible_utils.kv.put(b"test_key", b"test_val")
-        assert compatible_utils.kv.get(b"test_key") == b"test_val"
-
-        time.sleep(5)
-        fed.shutdown()
-        ray.shutdown()
-
-        assert compatible_utils.kv is None
-        with pytest.raises(ValueError):
-            # Make sure the kv actor is non-exist no matter whether it's in client mode
-            ray.get_actor("_INTERNAL_KV_ACTOR")
+def run(party):
+    compatible_utils.init_ray("local")
+    cluster = {
+        'alice': {'address': '127.0.0.1:11010', 'listen_addr': '0.0.0.0:11010'},
+        'bob': {'address': '127.0.0.1:11011', 'listen_addr': '0.0.0.0:11011'},
+    }
+    assert compatible_utils.kv is None
+    fed.init(cluster=cluster, party=party)
+    assert compatible_utils.kv
+    assert not compatible_utils.kv.put(b"test_key", b"test_val")
+    assert compatible_utils.kv.get(b"test_key") == b"test_val"
+
+    time.sleep(5)
+    fed.shutdown()
+
+    assert compatible_utils.kv is None
+    with pytest.raises(ValueError):
+        # Make sure the kv actor is non-exist no matter whether it's in client mode
+        ray.get_actor("_INTERNAL_KV_ACTOR")
+    ray.shutdown()
+
 
+def test_kv_init():
     p_alice = multiprocessing.Process(target=run, args=('alice',))
     p_bob = multiprocessing.Process(target=run, args=('bob',))
     p_alice.start()