fix: Update mock_redis to handle session compression in tests

- Add SmartRedisMock class that correctly handles compressed (:z suffix)
  and uncompressed session keys in tests
- Add set_uncompressed_return_value() for security tests on session data
- Add set_compressed_return_value() for testing compressed session retrieval
- Update test_redis_client_security.py to use new mock pattern
- Update test_get_session_compressed to use set_compressed_return_value()
- All 1142 tests now pass (was 10 failing due to mock not handling
  session compression logic correctly)

Author: H4LFdotDEV

python/tests/conftest.py

@@ -137,12 +137,83 @@ def sample_documents():
     ]
 
 
+class SmartRedisMock:
+    """Custom mock for Redis get() that handles session compression logic.
+
+    The redis_client tries compressed keys (ending in :z) first,
+    then falls back to uncompressed. This mock returns None for compressed
+    keys by default.
+
+    For session tests (uncompressed), use: mock.get.set_uncompressed_return_value(value)
+    For session tests (compressed), use: mock.get.set_compressed_return_value(value)
+    For other tests (template/query/embedding), use: mock.get.return_value = value
+    For error tests, use: mock.get.side_effect = Exception(...)
+    """
+
+    def __init__(self):
+        self._direct_value = None
+        self._session_value = None
+        self._compressed_value = None
+        self._side_effect = None
+        self.call_args_list = []
+
+    @property
+    def return_value(self):
+        return self._direct_value
+
+    @return_value.setter
+    def return_value(self, value):
+        self._direct_value = value
+
+    @property
+    def side_effect(self):
+        return self._side_effect
+
+    @side_effect.setter
+    def side_effect(self, value):
+        self._side_effect = value
+
+    def set_uncompressed_return_value(self, value):
+        """Set return value for uncompressed session keys."""
+        self._session_value = value
+
+    def set_compressed_return_value(self, value):
+        """Set return value for compressed session keys (:z suffix)."""
+        self._compressed_value = value
+
+    def __call__(self, key):
+        self.call_args_list.append((key,))
+
+        # If side_effect is set, raise/call it
+        if self._side_effect is not None:
+            if isinstance(self._side_effect, BaseException):
+                raise self._side_effect
+            elif callable(self._side_effect):
+                return self._side_effect(key)
+            else:
+                raise self._side_effect
+
+        # Compressed session keys (:z suffix)
+        if key.endswith(":z"):
+            if self._compressed_value is not None:
+                return self._compressed_value
+            return None
+        # Session keys use session value if set
+        if ":session:" in key and self._session_value is not None:
+            return self._session_value
+        # All other keys use direct value
+        return self._direct_value
+
+
 @pytest.fixture
 def mock_redis():
     """Mock Redis client for tests without Redis."""
     mock = MagicMock()
     mock.ping.return_value = True
-    mock.get.return_value = None
+
+    # Use smart mock for get() to handle session compression
+    mock.get = SmartRedisMock()
+
     mock.setex.return_value = True
     mock.delete.return_value = 1
     mock.keys.return_value = []

python/tests/test_redis_client.py

@@ -113,8 +113,8 @@ def test_get_session_compressed(self, redis_client, sample_session_state, mock_r
         compressed = zlib.compress(data.encode(), level=6)
         encoded = base64.b64encode(compressed).decode()
 
-        # Mock get to return compressed data on first call (compressed key)
-        mock_redis.get.return_value = encoded
+        # Mock get to return compressed data for compressed key (:z suffix)
+        mock_redis.get.set_compressed_return_value(encoded)
         result = redis_client.get_session(sample_session_state.session_id)
         assert result is not None
         assert result.session_id == sample_session_state.session_id

python/tests/test_redis_client_security.py

@@ -39,7 +39,8 @@ def redis_client(self, test_config, mock_redis):
 
     def test_get_session_valid_data(self, redis_client, sample_session_state, mock_redis):
         """Test get_session with valid data."""
-        mock_redis.get.return_value = json.dumps(sample_session_state.to_dict())
+        # Set return value for uncompressed key (compressed key returns None)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(sample_session_state.to_dict()))
 
         result = redis_client.get_session(sample_session_state.session_id)
 
@@ -60,15 +61,15 @@ def test_get_session_invalid_session_id_format(self, redis_client, mock_redis):
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(invalid_data)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(invalid_data))
 
         result = redis_client.get_session("test")
 
         assert result is None  # Validation failed
 
     def test_get_session_malformed_json(self, redis_client, mock_redis, caplog):
         """Test get_session handles malformed JSON gracefully."""
-        mock_redis.get.return_value = "{ invalid json"
+        mock_redis.get.set_uncompressed_return_value("{ invalid json")
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test-session")
@@ -83,7 +84,7 @@ def test_get_session_missing_required_field(self, redis_client, mock_redis, capl
             "project_path": "path",
             # Missing created_at, updated_at, etc.
         }
-        mock_redis.get.return_value = json.dumps(incomplete_data)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(incomplete_data))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test")
@@ -103,7 +104,7 @@ def test_get_session_invalid_timestamp(self, redis_client, mock_redis, caplog):
             "created_at": "invalid-timestamp",  # Invalid format
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(invalid_data)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(invalid_data))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test")
@@ -123,7 +124,7 @@ def test_get_session_path_traversal_attempt(self, redis_client, mock_redis, capl
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(traversal_data)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(traversal_data))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test")
@@ -143,7 +144,7 @@ def test_get_session_path_traversal_in_files(self, redis_client, mock_redis, cap
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(traversal_data)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(traversal_data))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test")
@@ -165,7 +166,7 @@ def test_get_session_extra_fields_rejected(self, redis_client, mock_redis, caplo
             "malicious_field": "should not be here",  # Extra field
             "another_injection": {"nested": "data"},
         }
-        mock_redis.get.return_value = json.dumps(data_with_extras)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(data_with_extras))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("test")
@@ -505,7 +506,7 @@ def test_session_id_redis_flush_injection(self, redis_client, mock_redis, caplog
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(redis_injection)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(redis_injection))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("whatever")
@@ -525,7 +526,7 @@ def test_session_id_shell_injection(self, redis_client, mock_redis, caplog):
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(shell_injection)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(shell_injection))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("whatever")
@@ -544,7 +545,7 @@ def test_session_id_backtick_injection(self, redis_client, mock_redis, caplog):
             "created_at": now,
             "updated_at": now,
         }
-        mock_redis.get.return_value = json.dumps(backtick_injection)
+        mock_redis.get.set_uncompressed_return_value(json.dumps(backtick_injection))
 
         with caplog.at_level(logging.ERROR):
             result = redis_client.get_session("whatever")
@@ -567,7 +568,7 @@ def test_redis_error_logging(self, redis_client, mock_redis, caplog):
 
     def test_session_state_maintains_type_safety(self, redis_client, sample_session_state, mock_redis):
         """Test SessionState object has correct types after deserialization."""
-        mock_redis.get.return_value = json.dumps(sample_session_state.to_dict())
+        mock_redis.get.set_uncompressed_return_value(json.dumps(sample_session_state.to_dict()))
 
         session = redis_client.get_session(sample_session_state.session_id)