From 33992eb398b8f007962633f305b86c0c41971f02 Mon Sep 17 00:00:00 2001
From: Guillermo Perez <bisho@revenuecat.com>
Date: Sun, 5 Nov 2023 22:02:22 +0100
Subject: [PATCH] Copy prebuilt default flags

---
 .../commands/high_level_commands.py           | 68 +++++++++++--------
 tests/probabilistic_hot_cache_test.py         |  2 +-
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/src/meta_memcache/commands/high_level_commands.py b/src/meta_memcache/commands/high_level_commands.py
index c5922b2..790b936 100644
--- a/src/meta_memcache/commands/high_level_commands.py
+++ b/src/meta_memcache/commands/high_level_commands.py
@@ -32,6 +32,22 @@
 T = TypeVar("T")
 _REFILL_FAILURE_HANDLING = FailureHandling(track_write_failures=False)
 
+DEFAULT_FLAGS: Set[Flag] = {
+    Flag.RETURN_VALUE,
+    Flag.RETURN_TTL,
+    Flag.RETURN_CLIENT_FLAG,
+    Flag.RETURN_LAST_ACCESS,
+    Flag.RETURN_FETCHED,
+}
+DEFAULT_CAS_FLAGS: Set[Flag] = {
+    Flag.RETURN_VALUE,
+    Flag.RETURN_TTL,
+    Flag.RETURN_CLIENT_FLAG,
+    Flag.RETURN_LAST_ACCESS,
+    Flag.RETURN_FETCHED,
+    Flag.RETURN_CAS_TOKEN,
+}
+
 
 class HighLevelCommandMixinWithMetaCommands(
     HighLevelCommandsProtocol, MetaCommandsProtocol, Protocol
@@ -330,20 +346,18 @@ def _multi_get(
         recache_policy: Optional[RecachePolicy] = None,
         return_cas_token: bool = False,
     ) -> Dict[Key, Optional[Value]]:
-        flags = {
-            Flag.RETURN_VALUE,
-            Flag.RETURN_TTL,
-            Flag.RETURN_CLIENT_FLAG,
-            Flag.RETURN_LAST_ACCESS,
-            Flag.RETURN_FETCHED,
-        }
         if return_cas_token:
-            flags.add(Flag.RETURN_CAS_TOKEN)
-        int_flags = {}
-        if recache_policy:
-            int_flags[IntFlag.RECACHE_TTL] = recache_policy.ttl
-        if touch_ttl is not None and touch_ttl >= 0:
-            int_flags[IntFlag.CACHE_TTL] = touch_ttl
+            flags = DEFAULT_CAS_FLAGS.copy()
+        else:
+            flags = DEFAULT_FLAGS.copy()
+        if recache_policy is None and touch_ttl is None:
+            int_flags = None
+        else:
+            int_flags = {}
+            if recache_policy:
+                int_flags[IntFlag.RECACHE_TTL] = recache_policy.ttl
+            if touch_ttl is not None and touch_ttl >= 0:
+                int_flags[IntFlag.CACHE_TTL] = touch_ttl
 
         results = self.meta_multiget(
             keys=[key if isinstance(key, Key) else Key(key) for key in keys],
@@ -379,22 +393,20 @@ def _get(
         return_cas_token: bool = False,
     ) -> Optional[Value]:
         key = key if isinstance(key, Key) else Key(key)
-        flags = {
-            Flag.RETURN_VALUE,
-            Flag.RETURN_TTL,
-            Flag.RETURN_CLIENT_FLAG,
-            Flag.RETURN_LAST_ACCESS,
-            Flag.RETURN_FETCHED,
-        }
         if return_cas_token:
-            flags.add(Flag.RETURN_CAS_TOKEN)
-        int_flags = {}
-        if lease_policy:
-            int_flags[IntFlag.MISS_LEASE_TTL] = lease_policy.ttl
-        if recache_policy:
-            int_flags[IntFlag.RECACHE_TTL] = recache_policy.ttl
-        if touch_ttl is not None and touch_ttl >= 0:
-            int_flags[IntFlag.CACHE_TTL] = touch_ttl
+            flags = DEFAULT_CAS_FLAGS.copy()
+        else:
+            flags = DEFAULT_FLAGS.copy()
+        if lease_policy is None and recache_policy is None and touch_ttl is None:
+            int_flags = None
+        else:
+            int_flags = {}
+            if lease_policy:
+                int_flags[IntFlag.MISS_LEASE_TTL] = lease_policy.ttl
+            if recache_policy:
+                int_flags[IntFlag.RECACHE_TTL] = recache_policy.ttl
+            if touch_ttl is not None and touch_ttl >= 0:
+                int_flags[IntFlag.CACHE_TTL] = touch_ttl
 
         result = self.meta_get(key, flags=flags, int_flags=int_flags)
         return self._process_get_result(key, result)
diff --git a/tests/probabilistic_hot_cache_test.py b/tests/probabilistic_hot_cache_test.py
index c1ce347..a87b0fb 100644
--- a/tests/probabilistic_hot_cache_test.py
+++ b/tests/probabilistic_hot_cache_test.py
@@ -85,7 +85,7 @@ def random(monkeypatch) -> Mock:
         Flag.RETURN_FETCHED,
         Flag.RETURN_CLIENT_FLAG,
     },
-    "int_flags": {},
+    "int_flags": None,
     "token_flags": None,
     "failure_handling": DEFAULT_FAILURE_HANDLING,
 }