Updated SegmentWorker

Author: chillaq

splitio/push/segmentworker.py

@@ -1,12 +1,29 @@
 """Segment changes processing worker."""
 import logging
 import threading
+import abc
+
+from splitio.optional.loaders import asyncio
 
 
 _LOGGER = logging.getLogger(__name__)
 
+class SegmentWorkerBase(object, metaclass=abc.ABCMeta):
+    """HttpClient wrapper template."""
+
+    @abc.abstractmethod
+    def is_running(self):
+        """Return whether the working is running."""
+
+    @abc.abstractmethod
+    def start(self):
+        """Start worker."""
 
-class SegmentWorker(object):
+    @abc.abstractmethod
+    def stop(self):
+        """Stop worker."""
+
+class SegmentWorker(SegmentWorkerBase):
     """Segment Worker for processing updates."""
 
     _centinel = object()
@@ -65,3 +82,61 @@ def stop(self):
             return
         self._running = False
         self._segment_queue.put(self._centinel)
+
+class SegmentWorkerAsync(SegmentWorkerBase):
+    """Segment Worker for processing updates."""
+
+    _centinel = object()
+
+    def __init__(self, synchronize_segment, segment_queue):
+        """
+        Class constructor.
+
+        :param synchronize_segment: handler to perform segment synchronization on incoming event
+        :type synchronize_segment: function
+
+        :param segment_queue: queue with segment updates notifications
+        :type segment_queue: asyncio.Queue
+        """
+        self._segment_queue = segment_queue
+        self._handler = synchronize_segment
+        self._running = False
+
+    def is_running(self):
+        """Return whether the working is running."""
+        return self._running
+
+    async def _run(self):
+        """Run worker handler."""
+        while self.is_running():
+            event = await self._segment_queue.get()
+            if not self.is_running():
+                break
+            if event == self._centinel:
+                continue
+            _LOGGER.debug('Processing segment_update: %s, change_number: %d',
+                          event.segment_name, event.change_number)
+            try:
+                await self._handler(event.segment_name, event.change_number)
+            except Exception:
+                _LOGGER.error('Exception raised in segment synchronization')
+                _LOGGER.debug('Exception information: ', exc_info=True)
+
+    def start(self):
+        """Start worker."""
+        if self.is_running():
+            _LOGGER.debug('Worker is already running')
+            return
+        self._running = True
+
+        _LOGGER.debug('Starting Segment Worker')
+        asyncio.get_event_loop().create_task(self._run())
+
+    async def stop(self):
+        """Stop worker."""
+        _LOGGER.debug('Stopping Segment Worker')
+        if not self.is_running():
+            _LOGGER.debug('Worker is not running. Ignoring.')
+            return
+        self._running = False
+        await self._segment_queue.put(self._centinel)

tests/push/test_segment_worker.py

@@ -4,8 +4,9 @@
 import pytest
 
 from splitio.api import APIException
-from splitio.push.segmentworker import SegmentWorker
+from splitio.push.segmentworker import SegmentWorker, SegmentWorkerAsync
 from splitio.models.notification import SegmentChangeNotification
+from splitio.optional.loaders import asyncio
 
 change_number_received = None
 segment_name_received = None
@@ -58,3 +59,52 @@ def test_handler(self):
 
         segment_worker.stop()
         assert not segment_worker.is_running()
+
+class SegmentWorkerAsyncTests(object):
+    async def test_on_error(self):
+        q = asyncio.Queue()
+
+        def handler_sync(change_number):
+            raise APIException('some')
+
+        segment_worker = SegmentWorkerAsync(handler_sync, q)
+        segment_worker.start()
+        assert segment_worker.is_running()
+
+        await q.put(SegmentChangeNotification('some', 'SEGMENT_UPDATE', 123456789, 'some'))
+
+        with pytest.raises(Exception):
+            segment_worker._handler()
+
+        assert segment_worker.is_running()
+        assert(self._worker_running())
+        await segment_worker.stop()
+        await asyncio.sleep(.1)
+        assert not segment_worker.is_running()
+        assert(not self._worker_running())
+
+    def _worker_running(self):
+        worker_running = False
+        for task in asyncio.Task.all_tasks():
+            if task._coro.cr_code.co_name == '_run' and not task.done():
+                worker_running = True
+                break
+        return worker_running
+
+    async def test_handler(self):
+        q = asyncio.Queue()
+        segment_worker = SegmentWorkerAsync(handler_sync, q)
+        global change_number_received
+        assert not segment_worker.is_running()
+        segment_worker.start()
+        assert segment_worker.is_running()
+
+        await q.put(SegmentChangeNotification('some', 'SEGMENT_UPDATE', 123456789, 'some'))
+
+        await asyncio.sleep(.1)
+        assert change_number_received == 123456789
+        assert segment_name_received == 'some'
+
+        await segment_worker.stop()
+        await asyncio.sleep(.1)
+        assert(not self._worker_running())