GG-33025: Add typing for Cluster and AioCluster (#39)

Author: isapego

pygridgain/aio_cluster.py

@@ -33,7 +33,7 @@ def __init__(self, client: 'AioClient'):
         self._client = client
 
     @status_to_exception(ClusterError)
-    async def get_state(self):
+    async def get_state(self) -> 'ClusterState':
         """
         Gets current cluster state.
 
@@ -43,7 +43,7 @@ async def get_state(self):
         return await cluster_get_state_async(await self._client.random_node())
 
     @status_to_exception(ClusterError)
-    async def set_state(self, state):
+    async def set_state(self, state: 'ClusterState'):
         """
         Changes current cluster state to the given.
 

pygridgain/api/cluster.py

@@ -16,7 +16,8 @@
 from pygridgain.api import APIResult
 from pygridgain.connection import AioConnection, Connection
 from pygridgain.datatypes import Byte
-from pygridgain.exceptions import NotSupportedByClusterError
+from pygridgain.datatypes.cluster_state import ClusterState
+from pygridgain.exceptions import NotSupportedByClusterError, ClusterError
 from pygridgain.queries import Query, query_perform
 from pygridgain.queries.op_codes import OP_CLUSTER_GET_STATE, OP_CLUSTER_CHANGE_STATE
 
@@ -44,7 +45,7 @@ async def cluster_get_state_async(connection: 'AioConnection', query_id=None) ->
 
 def __post_process_get_state(result):
     if result.status == 0:
-        result.value = result.value['state']
+        result.value = ClusterState(result.value['state'])
     return result
 
 
@@ -60,7 +61,7 @@ def __cluster_get_state(connection, query_id):
     )
 
 
-def cluster_set_state(connection: 'Connection', state: int, query_id=None) -> 'APIResult':
+def cluster_set_state(connection: 'Connection', state: ClusterState, query_id=None) -> 'APIResult':
     """
     Set cluster state.
 
@@ -75,7 +76,7 @@ def cluster_set_state(connection: 'Connection', state: int, query_id=None) -> 'A
     return __cluster_set_state(connection, state, query_id)
 
 
-async def cluster_set_state_async(connection: 'AioConnection', state: int, query_id=None) -> 'APIResult':
+async def cluster_set_state_async(connection: 'AioConnection', state: ClusterState, query_id=None) -> 'APIResult':
     """
     Async version of cluster_get_state
     """
@@ -89,6 +90,9 @@ def __post_process_set_state(result):
 
 
 def __cluster_set_state(connection, state, query_id):
+    if not ClusterState.has_value(state):
+        raise ClusterError(f'Unknown cluster state [state={state}]')
+
     if not connection.protocol_context.is_cluster_api_supported():
         raise NotSupportedByClusterError('Cluster API is not supported by the cluster')
 

pygridgain/api/key_value.py

@@ -57,7 +57,8 @@ def cache_put(connection: 'Connection', cache: Union[str, int], key: Any, value:
 
 
 async def cache_put_async(connection: 'AioConnection', cache: Union[str, int], key: Any, value: Any,
-                          key_hint: 'GridGainDataType' = None, value_hint: 'GridGainDataType' = None, binary: bool = False,
+                          key_hint: 'GridGainDataType' = None, value_hint: 'GridGainDataType' = None,
+                          binary: bool = False,
                           query_id: Optional[int] = None) -> 'APIResult':
     """
     Async version of cache_put
@@ -413,7 +414,8 @@ def __cache_get_and_put(connection, cache, key, value, key_hint, value_hint, bin
 
 
 def cache_get_and_replace(connection: 'Connection', cache: Union[str, int], key: Any, value: Any,
-                          key_hint: 'GridGainDataType' = None, value_hint: 'GridGainDataType' = None, binary: bool = False,
+                          key_hint: 'GridGainDataType' = None, value_hint: 'GridGainDataType' = None,
+                          binary: bool = False,
                           query_id: Optional[int] = None) -> 'APIResult':
     """
     Puts a value with a given key to cache, returning previous value
@@ -473,7 +475,8 @@ def __cache_get_and_replace(connection, cache, key, key_hint, value, value_hint,
     )
 
 
-def cache_get_and_remove(connection: 'Connection', cache: Union[str, int], key: Any, key_hint: 'GridGainDataType' = None,
+def cache_get_and_remove(connection: 'Connection', cache: Union[str, int], key: Any,
+                         key_hint: 'GridGainDataType' = None,
                          binary: bool = False, query_id: Optional[int] = None) -> 'APIResult':
     """
     Removes the cache entry with specified key, returning the value.
@@ -740,7 +743,8 @@ def cache_replace_if_equals(connection: 'Connection', cache: Union[str, int], ke
 
 async def cache_replace_if_equals_async(
         connection: 'AioConnection', cache: Union[str, int], key: Any, sample: Any, value: Any,
-        key_hint: 'GridGainDataType' = None, sample_hint: 'GridGainDataType' = None, value_hint: 'GridGainDataType' = None,
+        key_hint: 'GridGainDataType' = None, sample_hint: 'GridGainDataType' = None,
+        value_hint: 'GridGainDataType' = None,
         binary: bool = False, query_id: Optional[int] = None) -> 'APIResult':
     """
     Async version of cache_replace_if_equals.
@@ -1259,4 +1263,5 @@ def internal(result):
             result.value = result.value[key]
 
         return result
+
     return internal

pygridgain/cluster.py

@@ -34,7 +34,7 @@ def __init__(self, client: 'Client'):
         self._client = client
 
     @status_to_exception(ClusterError)
-    def get_state(self):
+    def get_state(self) -> 'ClusterState':
         """
         Gets current cluster state.
 
@@ -44,7 +44,7 @@ def get_state(self):
         return cluster_get_state(self._client.random_node)
 
     @status_to_exception(ClusterError)
-    def set_state(self, state):
+    def set_state(self, state: 'ClusterState'):
         """
         Changes current cluster state to the given.
 

pygridgain/connection/handshake.py

@@ -16,7 +16,7 @@
 from typing import Optional
 
 from pygridgain.connection.protocol_context import ProtocolContext
-from pygridgain.datatypes import Byte, ByteArrayObject, Int, MapObject, Short, String, UUIDObject, ByteArrayObject
+from pygridgain.datatypes import Byte, Int, MapObject, Short, String, UUIDObject, ByteArrayObject
 from pygridgain.datatypes.internal import Struct
 from pygridgain.stream import READ_BACKWARD
 

pygridgain/datatypes/cluster_state.py

@@ -27,3 +27,9 @@ class ClusterState(IntEnum):
     #: Cluster activated. Cache read operation allowed, Cache data change operation
     #: aren't allowed.
     ACTIVE_READ_ONLY = 2
+
+    @classmethod
+    def has_value(cls, value):
+        if not hasattr(cls, 'values'):
+            cls.values = set(item.value for item in ClusterState)
+        return value in cls.values

tests/affinity/test_affinity_bad_servers.py

@@ -20,6 +20,7 @@
 from tests.affinity.conftest import CLIENT_SOCKET_TIMEOUT
 from tests.util import start_ignite, kill_process_tree
 
+
 @pytest.fixture(params=['with-partition-awareness', 'without-partition-awareness'])
 def with_partition_awareness(request):
     yield request.param == 'with-partition-awareness'

tests/affinity/test_affinity_request_routing.py

@@ -27,7 +27,8 @@
 from pygridgain.constants import PROTOCOL_BYTE_ORDER
 from pygridgain.datatypes import String, LongObject
 from pygridgain.datatypes.cache_config import CacheMode
-from pygridgain.datatypes.prop_codes import PROP_NAME, PROP_BACKUPS_NUMBER, PROP_CACHE_KEY_CONFIGURATION, PROP_CACHE_MODE
+from pygridgain.datatypes.prop_codes import PROP_NAME, PROP_BACKUPS_NUMBER, PROP_CACHE_KEY_CONFIGURATION,\
+    PROP_CACHE_MODE
 from tests.util import wait_for_condition, wait_for_condition_async, start_ignite, kill_process_tree
 
 try:

tests/common/test_sql.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 #
 import pytest
-from pygridgain.exceptions import SQLError
 
 from pygridgain import AioClient
 from pygridgain.aio_cache import AioCache

tests/custom/test_cluster.py

@@ -17,7 +17,7 @@
 import pytest
 
 from pygridgain import Client, AioClient
-from pygridgain.exceptions import CacheError
+from pygridgain.exceptions import CacheError, ClusterError
 from tests.util import clear_ignite_work_dir, start_ignite_gen
 
 from pygridgain.datatypes.cluster_state import ClusterState
@@ -45,7 +45,17 @@ def server2(with_persistence, cleanup):
     yield from start_ignite_gen(idx=2, use_persistence=with_persistence)
 
 
-def test_cluster_set_active(with_persistence):
+def check_cluster_state_error(cluster, state: int):
+    with pytest.raises(ClusterError, match=f'Unknown cluster state \\[state={state}\\]'):
+        cluster.set_state(state)
+
+
+async def check_cluster_state_error_async(cluster, state: int):
+    with pytest.raises(ClusterError, match=f'Unknown cluster state \\[state={state}\\]'):
+        await cluster.set_state(state)
+
+
+def test_cluster_set_state(with_persistence):
     key = 42
     val = 42
     start_state = ClusterState.INACTIVE if with_persistence else ClusterState.ACTIVE
@@ -58,6 +68,11 @@ def test_cluster_set_active(with_persistence):
         cluster.set_state(ClusterState.ACTIVE)
         assert cluster.get_state() == ClusterState.ACTIVE
 
+        check_cluster_state_error(cluster, 3)
+        check_cluster_state_error(cluster, 42)
+        check_cluster_state_error(cluster, 1234567890987654321)
+        check_cluster_state_error(cluster, -1)
+
         cache = client.get_or_create_cache("test_cache")
         cache.put(key, val)
         assert cache.get(key) == val
@@ -86,7 +101,7 @@ def test_cluster_set_active(with_persistence):
 
 
 @pytest.mark.asyncio
-async def test_cluster_set_active_async(with_persistence):
+async def test_cluster_set_state_async(with_persistence):
     key = 42
     val = 42
     start_state = ClusterState.INACTIVE if with_persistence else ClusterState.ACTIVE
@@ -99,6 +114,11 @@ async def test_cluster_set_active_async(with_persistence):
         await cluster.set_state(ClusterState.ACTIVE)
         assert await cluster.get_state() == ClusterState.ACTIVE
 
+        await check_cluster_state_error_async(cluster, 3)
+        await check_cluster_state_error_async(cluster, 42)
+        await check_cluster_state_error_async(cluster, 1234567890987654321)
+        await check_cluster_state_error_async(cluster, -1)
+
         cache = await client.get_or_create_cache("test_cache")
         await cache.put(key, val)
         assert await cache.get(key) == val