fix and rewrite logic

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -1,20 +1,17 @@
 import asyncio
 import threading
-import time
 
 import pytest
 
 import ydb
 from ydb.aio.coordination import CoordinationClient as AioCoordinationClient
-from ydb import StatusCode, logger
+from ydb import StatusCode
 
 from ydb.coordination import (
     NodeConfig,
     ConsistencyMode,
     RateLimiterCountersMode,
     CoordinationClient,
-    CreateSemaphoreResult,
-    DescribeLockResult,
 )
 
 
@@ -120,165 +117,103 @@ async def test_coordination_node_lifecycle_async(self, async_coordination_node):
             await client.describe_node(node_path)
 
     async def test_coordination_lock_describe_full_async(self, async_coordination_node):
-        client, node_path, config = async_coordination_node
-
-        lock = client.lock("test_lock", node_path)
-
-        create = await lock.create(init_limit=1, init_data=b"hello")
-        assert create.status == StatusCode.SUCCESS
-
-        desc = await lock.describe()
-        assert desc.status == StatusCode.SUCCESS
-        assert desc.name == "test_lock"
-        assert desc.data == b"hello"
-        assert desc.count == 0
-        assert desc.ephemeral is False
-        assert list(desc.owners) == []
-        assert list(desc.waiters) == []
-
-        upd = await lock.update(new_data=b"world")
-        assert upd.status == StatusCode.SUCCESS
-
-        desc2 = await lock.describe()
-        assert desc2.status == StatusCode.SUCCESS
-        assert desc2.name == "test_lock"
-        assert desc2.data == b"world"
-        assert desc2.count == 0
-        assert desc2.ephemeral is False
-        assert list(desc2.owners) == []
-        assert list(desc2.waiters) == []
-
-        delete = await lock.delete()
-        assert delete.status == StatusCode.SUCCESS
-
-        desc_after = await lock.describe()
-        assert desc_after.status == StatusCode.NOT_FOUND
-
-    def test_coordination_lock_describe_full_sync(self, sync_coordination_node):
-        client, node_path, config = sync_coordination_node
-
-        lock = client.lock("test_lock", node_path)
-
-        create = lock.create(init_limit=1, init_data=b"hello")
-        assert create.status == StatusCode.SUCCESS
-
-        desc = lock.describe()
-        assert desc.status == StatusCode.SUCCESS
-        assert desc.name == "test_lock"
-        assert desc.data == b"hello"
-        assert desc.count == 0
-        assert desc.ephemeral is False
-        assert list(desc.owners) == []
-        assert list(desc.waiters) == []
-        upd = lock.update(new_data=b"world")
-        assert upd.status == StatusCode.SUCCESS
-
-        desc2 = lock.describe()
-        assert desc2.status == StatusCode.SUCCESS
-        assert desc2.name == "test_lock"
-        assert desc2.data == b"world"
-        assert desc2.count == 0
-        assert desc2.ephemeral is False
-        assert list(desc2.owners) == []
-        assert list(desc2.waiters) == []
-
-        delete = lock.delete()
-        assert delete.status == StatusCode.SUCCESS
-
-        desc_after = lock.describe()
-        assert desc_after.status == StatusCode.NOT_FOUND
+        client, node_path, _ = async_coordination_node
 
-    async def test_coordination_lock_racing_async(self, async_coordination_node):
-        client, node_path, initial_config = async_coordination_node
-        timeout = 5
+        async with client.node(node_path) as node:
+            lock = node.lock("test_lock")
 
-        lock = client.lock("test_lock", node_path)
-        await lock.create(init_limit=1, init_data=b"init-data")
+            desc = await lock.describe()
+            assert desc.status == StatusCode.NOT_FOUND
 
-        describe_resp: DescribeLockResult = await lock.describe()
-        assert describe_resp.status == StatusCode.SUCCESS
+            async with lock:
+                pass
 
-        lock2_started = asyncio.Event()
-        lock2_acquired = asyncio.Event()
-        lock2_release = asyncio.Event()
+            desc = await lock.describe()
+            assert desc.data == b""
 
-        async def second_lock_task():
-            lock2_started.set()
-            async with client.lock("test_lock", node_path):
-                lock2_acquired.set()
-                await lock2_release.wait()
+            await lock.update(new_data=b"world")
 
-        async with client.lock("test_lock", node_path) as lock1:
-            resp: DescribeLockResult = await lock1.describe()
-            assert resp.status == StatusCode.SUCCESS
+            desc2 = await lock.describe()
+            assert desc2.data == b"world"
 
-            t2 = asyncio.create_task(second_lock_task())
-            await asyncio.wait_for(lock2_started.wait(), timeout=timeout)
+    def test_coordination_lock_describe_full(self, sync_coordination_node):
+        client, node_path, _ = sync_coordination_node
 
-            await asyncio.sleep(0)
+        with client.node(node_path) as node:
+            lock = node.lock("test_lock")
 
-        await asyncio.wait_for(lock2_acquired.wait(), timeout=timeout)
+            desc = lock.describe()
+            assert desc.status == StatusCode.NOT_FOUND
 
-        lock2_release.set()
-        await asyncio.wait_for(t2, timeout=timeout)
+            with lock:
+                pass
 
-        delete_resp = await lock.delete()
-        assert delete_resp.status == StatusCode.SUCCESS
+            desc = lock.describe()
+            assert desc.data == b""
 
-        describe_after_delete: DescribeLockResult = await lock.describe()
-        assert describe_after_delete.status == StatusCode.NOT_FOUND
+            lock.update(new_data=b"world")
 
-    def test_coordination_lock_racing_sync(self, sync_coordination_node):
-        client, node_path, initial_config = sync_coordination_node
-        small_timeout = 1
+            desc2 = lock.describe()
+            assert desc2.data == b"world"
+
+    async def test_coordination_lock_racing_async(self, async_coordination_node):
+        client, node_path, _ = async_coordination_node
+        timeout = 5
 
-        lock = client.lock("test_lock", node_path)
+        async with client.node(node_path) as node:
+            lock2_started = asyncio.Event()
+            lock2_acquired = asyncio.Event()
+            lock2_release = asyncio.Event()
 
-        create_resp: CreateSemaphoreResult = lock.create(init_limit=1, init_data=b"init-data")
-        assert create_resp.status == StatusCode.SUCCESS
+            async def second_lock_task():
+                lock2_started.set()
+                async with node.lock("test_lock"):
+                    lock2_acquired.set()
+                    await lock2_release.wait()
 
-        describe_resp: DescribeLockResult = lock.describe()
-        assert describe_resp.status == StatusCode.SUCCESS
+            async with node.lock("test_lock"):
+                t2 = asyncio.create_task(second_lock_task())
+                await asyncio.wait_for(lock2_started.wait(), timeout=timeout)
 
-        lock2_ready = threading.Event()
-        lock2_acquired = threading.Event()
+            await asyncio.wait_for(lock2_acquired.wait(), timeout=timeout)
+            lock2_release.set()
+            await asyncio.wait_for(t2, timeout=timeout)
 
-        def second_lock_task():
-            try:
-                lock2_ready.set()
-                with client.lock("test_lock", node_path):
-                    lock2_acquired.set()
-                    logger.info("Second thread acquired lock")
-            except Exception as e:
-                logger.exception(f"{e} | second_lock_task failed")
+    def test_coordination_lock_racing(self, sync_coordination_node):
+        client, node_path, _ = sync_coordination_node
+        timeout = 5
+
+        with client.node(node_path) as node:
+            lock2_started = threading.Event()
+            lock2_acquired = threading.Event()
+            lock2_release = threading.Event()
 
-        t2 = threading.Thread(target=second_lock_task)
+            def second_lock_task():
+                lock2_started.set()
+                with node.lock("test_lock"):
+                    lock2_acquired.set()
+                    lock2_release.wait(timeout)
 
-        with client.lock("test_lock", node_path) as lock1:
-            resp = lock1.describe()
-            assert resp.status == StatusCode.SUCCESS
-            t2.start()
-            lock2_ready.wait(timeout=small_timeout)
+            with node.lock("test_lock"):
+                t2 = threading.Thread(target=second_lock_task)
+                t2.start()
 
-        lock2_acquired.wait(timeout=small_timeout)
-        t2.join(timeout=small_timeout)
+                assert lock2_started.wait(timeout)
 
-        delete_resp = lock.delete()
-        assert delete_resp.status == StatusCode.SUCCESS
-        time.sleep(small_timeout)
-        describe_after_delete: DescribeLockResult = lock.describe()
-        assert describe_after_delete.status == StatusCode.NOT_FOUND
+            assert lock2_acquired.wait(timeout)
+            lock2_release.set()
+            t2.join(timeout)
 
     async def test_coordination_reconnect_async(self, async_coordination_node):
-        client, node_path, config = async_coordination_node
+        client, node_path, _ = async_coordination_node
 
-        lock = client.lock("test_lock", node_path)
+        async with client.node(node_path) as node:
+            lock = node.lock("test_lock")
 
-        res = await lock.create(init_limit=1, init_data=b"")
-        assert res.status == StatusCode.SUCCESS
+            async with lock:
+                pass
 
-        await lock._reconnector._stream.close()
+            await node._reconnector.stop()
 
-        desc = await lock.describe()
-        assert desc.status == StatusCode.SUCCESS
-        assert desc.name == "test_lock"
+            async with lock:
+                pass

ydb/aio/coordination/client.py

@@ -8,8 +8,7 @@
 )
 from ..._grpc.grpcwrapper.ydb_coordination_public_types import NodeConfig
 from ...coordination.base import BaseCoordinationClient
-
-from .lock import CoordinationLock
+from .node import CoordinationNode
 
 
 class CoordinationClient(BaseCoordinationClient):
@@ -41,5 +40,5 @@ async def delete_node(self, path: str, settings=None):
             settings=settings,
         )
 
-    def lock(self, lock_name: str, node_path: str):
-        return CoordinationLock(self, lock_name, node_path=node_path)
+    def node(self, path: str) -> CoordinationNode:
+        return CoordinationNode(self._driver, path)