Merge pull request #16 from ATNoG/hotfix/sinkmanager_cancel_error

Finally fixed ig

Author: ThiagoAVicente

src/services/clickhouse.py

@@ -23,10 +23,34 @@ def get_data(self, batch_number: int = 1, batch_size: int = 50) -> list:
         return []
 
     def write_data(self, data: dict) -> None:
-        pass
+        """Write a single processed latency record to ClickHouse"""
+        try:
+            processed = ProcessedLatency(**data)
+            record = processed.to_dict()
+            
+            # Use async_insert for better performance
+            self.client.insert(
+                'analytics.processed_latency',
+                [record],
+                settings={'async_insert': 1, 'wait_for_async_insert': 0}
+            )
+        except Exception as e:
+            raise Exception(f"Failed to write to ClickHouse: {e}")
 
     def write_batch(self, data_list: list[dict]) -> None:
-        pass
+        """Write multiple processed latency records to ClickHouse"""
+        try:
+            processed_list = [ProcessedLatency(**d) for d in data_list]
+            records = [p.to_dict() for p in processed_list]
+            
+            # Use async_insert for better performance
+            self.client.insert(
+                'analytics.processed_latency',
+                records,
+                settings={'async_insert': 1, 'wait_for_async_insert': 0}
+            )
+        except Exception as e:
+            raise Exception(f"Failed to batch write to ClickHouse: {e}")
 
 
     def query_processed_latency(self,

src/services/influx.py

@@ -1,5 +1,5 @@
 from influxdb_client.client.influxdb_client import InfluxDBClient
-from influxdb_client.client.write_api import SYNCHRONOUS, WriteApi
+from influxdb_client.client.write_api import ASYNCHRONOUS, WriteApi
 from influxdb_client.client.query_api import QueryApi
 from src.configs.influx_conf import InfluxConf
 from src.models.raw import Raw
@@ -21,7 +21,7 @@ def connect(self):
         )
 
         # get iteration apis
-        self.write_api = self.client.write_api(write_options=SYNCHRONOUS)
+        self.write_api = self.client.write_api(write_options=ASYNCHRONOUS)
         self.query_api = self.client.query_api()
 
     def get_data(self, batch_number:int = 1 ,batch_size:int = 50) -> list[Raw]:

src/sink.py

@@ -60,6 +60,11 @@ async def start(self, *topics):
         self.bridge.add_n_topics(topics, bind=self.route_message)
 
         await self.bridge.start_consumer()
+        
+        # Keep the event loop alive - wait for the consumer task to complete
+        # This prevents asyncio.run() from exiting and cancelling the consumer task
+        if self.bridge._consumer_task:
+            await self.bridge._consumer_task
 
     async def stop(self):
         if self.bridge is not None:

tests/conftest.py

@@ -12,6 +12,7 @@ def __init__(self, *topics, hostname=None, port=None):
         self.port = port
         self._closed = False
         self._callbacks = {}
+        self._consumer_task = None  # Added for consumer task tracking
 
     def add_n_topics(self, topics, bind=None):
         """Mock add_n_topics method."""

tests/test_clickhouse_service.py

@@ -39,6 +39,59 @@ def test_connect(self, mock_clickhouse_client):
 
         assert service.client == mock_clickhouse_client
 
+    def test_write_data(self, clickhouse_service, mock_clickhouse_client):
+        """Test writing a single ProcessedLatency record to ClickHouse."""
+        # Prepare test data
+        test_data = {
+            'window_start_time': datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+            'window_end_time': datetime(2024, 1, 1, 12, 5, 0, tzinfo=timezone.utc),
+            'window_duration_seconds': 300.0,
+            'cell_index': 1,
+            'network': '5G',
+            'rsrp_mean': -80.0,
+            'rsrp_max': -70.0,
+            'rsrp_min': -90.0,
+            'rsrp_std': 5.0,
+            'sinr_mean': 15.0,
+            'sinr_max': 20.0,
+            'sinr_min': 10.0,
+            'sinr_std': 3.0,
+            'rsrq_mean': -10.0,
+            'rsrq_max': -8.0,
+            'rsrq_min': -12.0,
+            'rsrq_std': 1.5,
+            'latency_mean': 20.0,
+            'latency_max': 30.0,
+            'latency_min': 10.0,
+            'latency_std': 5.0,
+            'cqi_mean': 12.0,
+            'cqi_max': 15.0,
+            'cqi_min': 10.0,
+            'cqi_std': 2.0,
+            'primary_bandwidth': 100.0,
+            'ul_bandwidth': 50.0,
+            'sample_count': 100
+        }
+
+        # Execute write
+        clickhouse_service.write_data(test_data)
+
+        # Verify insert was called with correct parameters
+        mock_clickhouse_client.insert.assert_called_once()
+        call_args = mock_clickhouse_client.insert.call_args
+
+        # Check table name
+        assert call_args[0][0] == 'analytics.processed_latency'
+        
+        # Check data format (should be a list with one dict)
+        assert isinstance(call_args[0][1], list)
+        assert len(call_args[0][1]) == 1
+        
+        # Check async insert settings
+        assert call_args[1]['settings']['async_insert'] == 1
+        assert call_args[1]['settings']['wait_for_async_insert'] == 0
+
+
     def test_query_processed_latency_success(self, clickhouse_service, mock_clickhouse_client):
         """Test successful query of processed latency data."""
         # Mock query result