Fix unit tests

Author: incrypto32

tests/unit/test_streaming_helpers.py

@@ -217,9 +217,8 @@ class TestProcessBatchNonTransactional:
     """Test _process_batch_non_transactional helper method"""
 
     def test_successful_non_transactional_load(self, mock_loader, sample_batch, sample_ranges):
-        """Test successful non-transactional batch load"""
-        # Setup - mock state store for new unified system
-        mock_loader.state_store.is_processed = Mock(return_value=False)
+        """Test successful non-transactional batch load with watermark"""
+        # Setup - mock state store for watermark checkpointing
         mock_loader.state_store.mark_processed = Mock()
 
         # Mock load_batch to return success
@@ -228,86 +227,94 @@ def test_successful_non_transactional_load(self, mock_loader, sample_batch, samp
         )
         mock_loader.load_batch = Mock(return_value=success_result)
 
-        # Execute
+        # Execute with ranges_complete=True (watermark)
         result = mock_loader._process_batch_non_transactional(
             batch_data=sample_batch,
             table_name='test_table',
             connection_name='test_conn',
             ranges=sample_ranges,
-            batch_hash='hash123',
+            ranges_complete=True,
         )
 
         # Verify
         assert result.success
         assert result.rows_loaded == 3
 
-        # Verify method calls with state store
-        mock_loader.state_store.is_processed.assert_called_once()
+        # Verify method calls
         mock_loader.load_batch.assert_called_once()
+        # mark_processed should be called for watermarks (ranges_complete=True)
         mock_loader.state_store.mark_processed.assert_called_once()
 
-    def test_duplicate_detection_returns_skip_result(self, mock_loader, sample_batch, sample_ranges):
-        """Test duplicate detection returns skip result"""
-        # Setup - is_processed returns True
-        mock_loader.state_store.is_processed = Mock(return_value=True)
-        mock_loader.load_batch = Mock()  # Should not be called
+    def test_data_batch_does_not_checkpoint(self, mock_loader, sample_batch, sample_ranges):
+        """Test that data batches (ranges_complete=False) don't checkpoint state"""
+        # Setup
+        mock_loader.state_store.mark_processed = Mock()
 
-        # Execute
+        # Mock load_batch to return success
+        success_result = LoadResult(
+            rows_loaded=3, duration=0.1, ops_per_second=30.0, table_name='test_table', loader_type='mock', success=True
+        )
+        mock_loader.load_batch = Mock(return_value=success_result)
+
+        # Execute with ranges_complete=False (data batch)
         result = mock_loader._process_batch_non_transactional(
             batch_data=sample_batch,
             table_name='test_table',
             connection_name='test_conn',
             ranges=sample_ranges,
-            batch_hash='hash123',
+            ranges_complete=False,
         )
 
         # Verify
         assert result.success
-        assert result.rows_loaded == 0
-        assert result.metadata['operation'] == 'skip_duplicate'
-        assert result.metadata['ranges'] == [r.to_dict() for r in sample_ranges]
+        assert result.rows_loaded == 3
 
-        # load_batch should not be called for duplicates
-        mock_loader.load_batch.assert_not_called()
+        # load_batch should be called
+        mock_loader.load_batch.assert_called_once()
+        # mark_processed should NOT be called for data batches (ranges_complete=False)
+        mock_loader.state_store.mark_processed.assert_not_called()
 
-    def test_no_ranges_skips_duplicate_check(self, mock_loader, sample_batch):
-        """Test that no ranges means no duplicate checking"""
+    def test_no_ranges_skips_checkpoint(self, mock_loader, sample_batch):
+        """Test that no ranges means no checkpointing"""
         # Setup
-        mock_loader.state_store.is_processed = Mock()
+        mock_loader.state_store.mark_processed = Mock()
         success_result = LoadResult(
             rows_loaded=3, duration=0.1, ops_per_second=30.0, table_name='test_table', loader_type='mock', success=True
         )
         mock_loader.load_batch = Mock(return_value=success_result)
 
         # Execute with None ranges
         result = mock_loader._process_batch_non_transactional(
-            batch_data=sample_batch, table_name='test_table', connection_name='test_conn', ranges=None, batch_hash=None
+            batch_data=sample_batch,
+            table_name='test_table',
+            connection_name='test_conn',
+            ranges=None,
+            ranges_complete=True,
         )
 
         # Verify
         assert result.success
 
-        # is_processed should not be called
-        mock_loader.state_store.is_processed.assert_not_called()
+        # mark_processed should not be called when ranges is None
+        mock_loader.state_store.mark_processed.assert_not_called()
 
     def test_mark_processed_failure_continues(self, mock_loader, sample_batch, sample_ranges):
         """Test that mark_processed failure doesn't fail the load"""
         # Setup
-        mock_loader.state_store.is_processed = Mock(return_value=False)
         mock_loader.state_store.mark_processed = Mock(side_effect=Exception('Mark failed'))
 
         success_result = LoadResult(
             rows_loaded=3, duration=0.1, ops_per_second=30.0, table_name='test_table', loader_type='mock', success=True
         )
         mock_loader.load_batch = Mock(return_value=success_result)
 
-        # Execute - should not raise exception
+        # Execute - should not raise exception (watermark checkpoint failure is logged but not fatal)
         result = mock_loader._process_batch_non_transactional(
             batch_data=sample_batch,
             table_name='test_table',
             connection_name='test_conn',
             ranges=sample_ranges,
-            batch_hash='hash123',
+            ranges_complete=True,
         )
 
         # Verify - load still succeeded despite mark_processed failure