IGNITE-14429 Fix cache.get_size with non-default PeekModes

This closes #24

Author: ivandasch

pyignite/aio_cache.py

@@ -572,13 +572,13 @@ async def replace_if_equals(self, key, sample, value, key_hint=None, sample_hint
         return result
 
     @status_to_exception(CacheError)
-    async def get_size(self, peek_modes=0):
+    async def get_size(self, peek_modes=None):
         """
         Gets the number of entries in cache.
 
         :param peek_modes: (optional) limit count to near cache partition
          (PeekModes.NEAR), primary cache (PeekModes.PRIMARY), or backup cache
-         (PeekModes.BACKUP). Defaults to all cache partitions (PeekModes.ALL),
+         (PeekModes.BACKUP). Defaults to primary cache partitions (PeekModes.PRIMARY),
         :return: integer number of cache entries.
         """
         conn = await self.get_best_node()

pyignite/api/key_value.py

@@ -23,9 +23,8 @@
     OP_CACHE_CLEAR_KEYS, OP_CACHE_REMOVE_KEY, OP_CACHE_REMOVE_IF_EQUALS, OP_CACHE_REMOVE_KEYS, OP_CACHE_REMOVE_ALL,
     OP_CACHE_GET_SIZE, OP_CACHE_LOCAL_PEEK
 )
-from pyignite.datatypes import Map, Bool, Byte, Int, Long, AnyDataArray, AnyDataObject
+from pyignite.datatypes import Map, Bool, Byte, Int, Long, AnyDataArray, AnyDataObject, ByteArray
 from pyignite.datatypes.base import IgniteDataType
-from pyignite.datatypes.key_value import PeekModes
 from pyignite.queries import Query, query_perform
 from pyignite.utils import cache_id
 
@@ -1128,7 +1127,7 @@ def __cache_remove_all(connection, cache, binary, query_id):
     )
 
 
-def cache_get_size(connection: 'Connection', cache: Union[str, int], peek_modes: Union[int, list, tuple] = 0,
+def cache_get_size(connection: 'Connection', cache: Union[str, int], peek_modes: Union[int, list, tuple] = None,
                    binary: bool = False, query_id: Optional[int] = None) -> 'APIResult':
     """
     Gets the number of entries in cache.
@@ -1137,7 +1136,7 @@ def cache_get_size(connection: 'Connection', cache: Union[str, int], peek_modes:
     :param cache: name or ID of the cache,
     :param peek_modes: (optional) limit count to near cache partition
      (PeekModes.NEAR), primary cache (PeekModes.PRIMARY), or backup cache
-     (PeekModes.BACKUP). Defaults to all cache partitions (PeekModes.ALL),
+     (PeekModes.BACKUP). Defaults to pimary cache partitions (PeekModes.PRIMARY),
     :param binary: (optional) pass True to keep the value in binary form.
      False by default,
     :param query_id: (optional) a value generated by client and returned as-is
@@ -1151,21 +1150,23 @@ def cache_get_size(connection: 'Connection', cache: Union[str, int], peek_modes:
 
 
 async def cache_get_size_async(connection: 'AioConnection', cache: Union[str, int],
-                               peek_modes: Union[int, list, tuple] = 0, binary: bool = False,
+                               peek_modes: Union[int, list, tuple] = None, binary: bool = False,
                                query_id: Optional[int] = None) -> 'APIResult':
     return await __cache_get_size(connection, cache, peek_modes, binary, query_id)
 
 
 def __cache_get_size(connection, cache, peek_modes, binary, query_id):
-    if not isinstance(peek_modes, (list, tuple)):
-        peek_modes = [peek_modes] if peek_modes else []
+    if peek_modes is None:
+        peek_modes = []
+    elif not isinstance(peek_modes, (list, tuple)):
+        peek_modes = [peek_modes]
 
     query_struct = Query(
         OP_CACHE_GET_SIZE,
         [
             ('hash_code', Int),
             ('flag', Byte),
-            ('peek_modes', PeekModes),
+            ('peek_modes', ByteArray),
         ],
         query_id=query_id,
     )
@@ -1184,7 +1185,7 @@ def __cache_get_size(connection, cache, peek_modes, binary, query_id):
 
 
 def cache_local_peek(conn: 'Connection', cache: Union[str, int], key: Any, key_hint: 'IgniteDataType' = None,
-                     peek_modes: Union[int, list, tuple] = 0, binary: bool = False,
+                     peek_modes: Union[int, list, tuple] = None, binary: bool = False,
                      query_id: Optional[int] = None) -> 'APIResult':
     """
     Peeks at in-memory cached value using default optional peek mode.
@@ -1199,7 +1200,7 @@ def cache_local_peek(conn: 'Connection', cache: Union[str, int], key: Any, key_h
      should be converted,
     :param peek_modes: (optional) limit count to near cache partition
      (PeekModes.NEAR), primary cache (PeekModes.PRIMARY), or backup cache
-     (PeekModes.BACKUP). Defaults to all cache partitions (PeekModes.ALL),
+     (PeekModes.BACKUP). Defaults to primary cache partitions (PeekModes.PRIMARY),
     :param binary: (optional) pass True to keep the value in binary form.
      False by default,
     :param query_id: (optional) a value generated by client and returned as-is
@@ -1213,24 +1214,27 @@ def cache_local_peek(conn: 'Connection', cache: Union[str, int], key: Any, key_h
 
 async def cache_local_peek_async(
         conn: 'AioConnection', cache: Union[str, int], key: Any, key_hint: 'IgniteDataType' = None,
-        peek_modes: Union[int, list, tuple] = 0, binary: bool = False, query_id: Optional[int] = None) -> 'APIResult':
+        peek_modes: Union[int, list, tuple] = None, binary: bool = False,
+        query_id: Optional[int] = None) -> 'APIResult':
     """
     Async version of cache_local_peek.
     """
     return await __cache_local_peek(conn, cache, key, key_hint, peek_modes, binary, query_id)
 
 
 def __cache_local_peek(conn, cache, key, key_hint, peek_modes, binary, query_id):
-    if not isinstance(peek_modes, (list, tuple)):
-        peek_modes = [peek_modes] if peek_modes else []
+    if peek_modes is None:
+        peek_modes = []
+    elif not isinstance(peek_modes, (list, tuple)):
+        peek_modes = [peek_modes]
 
     query_struct = Query(
         OP_CACHE_LOCAL_PEEK,
         [
             ('hash_code', Int),
             ('flag', Byte),
             ('key', key_hint or AnyDataObject),
-            ('peek_modes', PeekModes),
+            ('peek_modes', ByteArray),
         ],
         query_id=query_id,
     )

pyignite/cache.py

@@ -694,13 +694,13 @@ def replace_if_equals(
         return result
 
     @status_to_exception(CacheError)
-    def get_size(self, peek_modes=0):
+    def get_size(self, peek_modes=None):
         """
         Gets the number of entries in cache.
 
         :param peek_modes: (optional) limit count to near cache partition
          (PeekModes.NEAR), primary cache (PeekModes.PRIMARY), or backup cache
-         (PeekModes.BACKUP). Defaults to all cache partitions (PeekModes.ALL),
+         (PeekModes.BACKUP). Defaults to primary cache partitions (PeekModes.PRIMARY),
         :return: integer number of cache entries.
         """
         return cache_get_size(

pyignite/datatypes/key_value.py

@@ -13,14 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .primitive_arrays import ByteArray
+from enum import IntEnum
 
 
-class PeekModes(ByteArray):
-
-    ALL = 1
-    NEAR = 2
-    PRIMARY = 4
-    BACKUP = 8
-    ONHEAP = 16
-    OFFHEAP = 32
+class PeekModes(IntEnum):
+    ALL = 0
+    NEAR = 1
+    PRIMARY = 2
+    BACKUP = 3
+    ONHEAP = 4
+    OFFHEAP = 5

tests/common/test_cache_size.py

@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from pyignite.datatypes.key_value import PeekModes
+from pyignite.datatypes.prop_codes import PROP_NAME, PROP_IS_ONHEAP_CACHE_ENABLED, PROP_BACKUPS_NUMBER
+from tests.util import get_or_create_cache, get_or_create_cache_async
+
+test_params = [
+    [
+        {
+            PROP_NAME: 'cache_onheap_backups_2',
+            PROP_IS_ONHEAP_CACHE_ENABLED: True,
+            PROP_BACKUPS_NUMBER: 2
+        },
+        [
+            [None, 1],
+            [PeekModes.PRIMARY, 1],
+            [PeekModes.BACKUP, 2],
+            [PeekModes.ALL, 3],
+            [[PeekModes.PRIMARY, PeekModes.BACKUP], 3],
+            [PeekModes.ONHEAP, 1],
+            [PeekModes.OFFHEAP, 1]
+        ]
+    ]
+]
+
+
+@pytest.mark.parametrize("cache_settings, cache_sizes", test_params)
+def test_cache_size(client, cache_settings, cache_sizes):
+    with get_or_create_cache(client, cache_settings) as cache:
+        cache.put(1, 1)
+
+        for props, exp_value in cache_sizes:
+            value = cache.get_size(props)
+            assert value == exp_value, f"expected {exp_value} for {props}, got {value} instead."
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cache_settings, cache_sizes", test_params)
+async def test_cache_size_async(async_client, cache_settings, cache_sizes):
+    async with get_or_create_cache_async(async_client, cache_settings) as cache:
+        await cache.put(1, 1)
+
+        for props, exp_value in cache_sizes:
+            value = await cache.get_size(props)
+            assert value == exp_value, f"expected {exp_value} for {props}, got {value} instead."

tests/util.py

@@ -35,17 +35,17 @@
 
 
 @contextlib.contextmanager
-def get_or_create_cache(client, cache_name):
-    cache = client.get_or_create_cache(cache_name)
+def get_or_create_cache(client, settings):
+    cache = client.get_or_create_cache(settings)
     try:
         yield cache
     finally:
         cache.destroy()
 
 
 @asynccontextmanager
-async def get_or_create_cache_async(client, cache_name):
-    cache = await client.get_or_create_cache(cache_name)
+async def get_or_create_cache_async(client, settings):
+    cache = await client.get_or_create_cache(settings)
     try:
         yield cache
     finally: