Use mock_message_context

Author: gwaramadze

quixstreams/dataframe/windows/time_based.py

@@ -581,17 +581,7 @@ def _on_expired_session(
         timestamp_ms: int,
         late_by_ms: int,
     ) -> None:
-        try:
-            ctx = message_context()
-            topic = ctx.topic
-            partition = ctx.partition
-            offset = ctx.offset
-        except:
-            # In test environments, message context might not be available
-            topic = "unknown"
-            partition = -1
-            offset = -1
-
+        ctx = message_context()
         to_log = True
 
         # Trigger the "on_late" callback if provided
@@ -604,9 +594,9 @@ def _on_expired_session(
                 start,
                 end,
                 self._name,
-                topic,
-                partition,
-                offset,
+                ctx.topic,
+                ctx.partition,
+                ctx.offset,
             )
         if to_log:
             logger.warning(
@@ -615,8 +605,8 @@ def _on_expired_session(
                 f"session={(start, end)} "
                 f"late_by_ms={late_by_ms} "
                 f"store_name={self._name} "
-                f"partition={topic}[{partition}] "
-                f"offset={offset}"
+                f"partition={ctx.topic}[{ctx.partition}] "
+                f"offset={ctx.offset}"
             )
 
 

tests/test_quixstreams/test_dataframe/fixtures.py

@@ -1,5 +1,5 @@
 from typing import Optional
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -59,3 +59,9 @@ def factory(
         return sdf
 
     return factory
+
+
+@pytest.fixture
+def mock_message_context():
+    with patch("quixstreams.dataframe.windows.time_based.message_context"):
+        yield

tests/test_quixstreams/test_dataframe/test_windows/test_session.py

@@ -549,7 +549,7 @@ def test_session_partition_expiration(
             assert key2 in expired_keys
 
     def test_session_window_late_events(
-        self, session_window_definition_factory, state_manager
+        self, session_window_definition_factory, state_manager, mock_message_context
     ):
         """Test handling of late events that arrive after session closure"""
         window_def = session_window_definition_factory(

tests/test_quixstreams/test_dataframe/test_windows/test_sliding.py

@@ -2,7 +2,6 @@
 from dataclasses import dataclass, field
 from itertools import chain
 from typing import Any
-from unittest import mock
 
 import pytest
 
@@ -741,12 +740,6 @@ def expected_windows_in_state(self) -> set[tuple[int, int]]:
 ]
 
 
-@pytest.fixture
-def mock_message_context():
-    with mock.patch("quixstreams.dataframe.windows.time_based.message_context"):
-        yield
-
-
 @pytest.fixture
 def sliding_window_definition_factory(
     state_manager, dataframe_factory, topic_manager_topic_factory