Repository: liampauling/flumine
Pull request #166 (2020-05-21 07:49:47)

worker refactor so function takes context and flumine only by default


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 081eded..a37d905 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,13 @@
 Release History
 ---------------
 
+1.5.1 (2020-05-21)
++++++++++++++++++++
+
+**Improvements**
+
+- Worker refactor to make init simpler when adding custom workers
+
 1.5.0 (2020-05-21)
 +++++++++++++++++++
 
diff --git a/flumine/__version__.py b/flumine/__version__.py
index 84895d8..7392a80 100644
--- a/flumine/__version__.py
+++ b/flumine/__version__.py
@@ -1,6 +1,6 @@
 __title__ = "flumine"
 __description__ = "Betfair trading framework"
 __url__ = "https://github.com/liampauling/flumine"
-__version__ = "1.5.0"
+__version__ = "1.5.1"
 __author__ = "Liam Pauling"
 __license__ = "MIT"
diff --git a/flumine/flumine.py b/flumine/flumine.py
index d053d74..9fb5af0 100644
--- a/flumine/flumine.py
+++ b/flumine/flumine.py
@@ -55,36 +55,33 @@ class Flumine(BaseFlumine):
     def _add_default_workers(self):
         self.add_worker(
             worker.BackgroundWorker(
-                interval=1200,
-                function=worker.keep_alive,
-                name="keep_alive",
-                func_args=(self.client,),
+                self, function=worker.keep_alive, interval=1200, name="keep_alive",
             )
         )
         self.add_worker(
             worker.BackgroundWorker(
-                start_delay=5,  # wait for login
-                interval=60,
+                self,
                 function=worker.poll_account_balance,
+                interval=60,
+                start_delay=5,  # wait for login
                 name="poll_account_balance",
-                func_args=(self, self.client),
             )
         )
         self.add_worker(
             worker.BackgroundWorker(
-                start_delay=5,  # wait for streams to populate
-                interval=60,
+                self,
                 function=worker.poll_market_catalogue,
+                interval=60,
+                start_delay=5,  # wait for streams to populate
                 name="poll_market_catalogue",
-                func_args=(self.client, self.markets, self.handler_queue),
             )
         )
         self.add_worker(
             worker.BackgroundWorker(
-                interval=10,  # restart
+                self,
                 function=worker.poll_cleared_orders,
+                interval=10,  # restart
                 name="poll_cleared_orders",
-                func_args=(self, self.client),
             )
         )
 
diff --git a/flumine/worker.py b/flumine/worker.py
index 4c1a859..10473ab 100644
--- a/flumine/worker.py
+++ b/flumine/worker.py
@@ -14,8 +14,9 @@ logger = logging.getLogger(__name__)
 class BackgroundWorker(threading.Thread):
     def __init__(
         self,
-        interval: int,
+        flumine,
         function,
+        interval: int,
         func_args: tuple = None,
         func_kwargs: dict = None,
         start_delay: int = 0,
@@ -23,8 +24,9 @@ class BackgroundWorker(threading.Thread):
         **kwargs
     ):
         threading.Thread.__init__(self, daemon=True, **kwargs)
-        self.interval = interval
+        self.flumine = flumine
         self.function = function
+        self.interval = interval
         self.func_args = func_args if func_args is not None else []
         self.func_kwargs = func_kwargs if func_kwargs is not None else {}
         self.start_delay = start_delay
@@ -50,7 +52,9 @@ class BackgroundWorker(threading.Thread):
                 },
             )
             try:
-                self.function(self.context, *self.func_args, **self.func_kwargs)
+                self.function(
+                    self.context, self.flumine, *self.func_args, **self.func_kwargs
+                )
             except Exception as e:
                 logger.error(
                     "Error in BackgroundWorker {0}: {1}".format(self.name, e),
@@ -63,10 +67,11 @@ class BackgroundWorker(threading.Thread):
             time.sleep(self.interval)
 
 
-def keep_alive(context: dict, client) -> None:
+def keep_alive(context: dict, flumine) -> None:
     """ Attempt keep alive if required or
     login if keep alive failed
     """
+    client = flumine.client
     if client.betting_client.session_token:
         try:
             resp = client.keep_alive()
@@ -89,10 +94,9 @@ def keep_alive(context: dict, client) -> None:
         )
 
 
-def poll_market_catalogue(
-    context: dict, client, markets, handler_queue: queue.Queue
-) -> None:
-    live_markets = list(markets.markets.keys())
+def poll_market_catalogue(context: dict, flumine) -> None:
+    client = flumine.client
+    live_markets = list(flumine.markets.markets.keys())
     for market_ids in chunks(live_markets, 100):
         try:
             market_catalogues = client.betting_client.betting.list_market_catalogue(
@@ -117,16 +121,18 @@ def poll_market_catalogue(
             continue
 
         if market_catalogues:
-            handler_queue.put(events.MarketCatalogueEvent(market_catalogues))
+            flumine.handler_queue.put(events.MarketCatalogueEvent(market_catalogues))
 
 
-def poll_account_balance(context: dict, flumine, client) -> None:
+def poll_account_balance(context: dict, flumine) -> None:
+    client = flumine.client
     client.update_account_details()
     if client.account_funds:
         flumine.log_control(events.BalanceEvent(client.account_funds))
 
 
-def poll_cleared_orders(context: dict, flumine, client) -> None:
+def poll_cleared_orders(context: dict, flumine) -> None:
+    client = flumine.client
     processed_orders, processed_markets, market_cache = [], [], []
     while True:
         market_id = flumine.cleared_market_queue.get()


Test patch:
diff --git a/tests/test_worker.py b/tests/test_worker.py
index 480622d..bab95e7 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -9,13 +9,21 @@ from flumine import worker
 class BackgroundWorkerTest(unittest.TestCase):
     def setUp(self):
         self.mock_function = mock.Mock()
+        self.mock_flumine = mock.Mock()
         self.worker = worker.BackgroundWorker(
-            123, self.mock_function, (1, 2), {"hello": "world"}, 5, {1: 2}
+            self.mock_flumine,
+            self.mock_function,
+            123,
+            (1, 2),
+            {"hello": "world"},
+            5,
+            {1: 2},
         )
 
     def test_init(self):
         self.assertEqual(self.worker.interval, 123)
         self.assertEqual(self.worker.function, self.mock_function)
+        self.assertEqual(self.worker.flumine, self.mock_flumine)
         self.assertEqual(self.worker.func_args, (1, 2))
         self.assertEqual(self.worker.func_kwargs, {"hello": "world"})
         self.assertEqual(self.worker.start_delay, 5)
@@ -31,56 +39,52 @@ class WorkersTest(unittest.TestCase):
 
     def test_keep_alive(self):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_client.betting_client.session_token = None
-        worker.keep_alive(mock_context, mock_client)
-        mock_client.login.assert_called_with()
+        mock_flumine = mock.Mock()
+        mock_flumine.client.betting_client.session_token = None
+        worker.keep_alive(mock_context, mock_flumine)
+        mock_flumine.client.login.assert_called_with()
 
-        mock_client.betting_client.session_token = 1
-        mock_client.betting_client.session_expired = True
-        worker.keep_alive(mock_context, mock_client)
-        mock_client.keep_alive.assert_called_with()
+        mock_flumine.client.betting_client.session_token = 1
+        mock_flumine.client.betting_client.session_expired = True
+        worker.keep_alive(mock_context, mock_flumine)
+        mock_flumine.client.keep_alive.assert_called_with()
 
     def test_keep_alive_failure(self):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_client.betting_client.session_token = None
+        mock_flumine = mock.Mock()
+        mock_flumine.client.betting_client.session_token = None
         mock_response = mock.Mock()
         mock_response.status = "FAILURE"
-        mock_client.betting_client.keep_alive.return_value = mock_response
-        worker.keep_alive(mock_context, mock_client)
-        mock_client.login.assert_called_with()
+        mock_flumine.client.betting_client.keep_alive.return_value = mock_response
+        worker.keep_alive(mock_context, mock_flumine)
+        mock_flumine.client.login.assert_called_with()
 
     def test_keep_alive_error(self):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_client.betting_client.session_token = None
-        mock_client.betting_client.keep_alive.side_effect = BetfairError
-        worker.keep_alive(mock_context, mock_client)
-        mock_client.login.assert_called_with()
+        mock_flumine = mock.Mock()
+        mock_flumine.client.betting_client.session_token = None
+        mock_flumine.client.betting_client.keep_alive.side_effect = BetfairError
+        worker.keep_alive(mock_context, mock_flumine)
+        mock_flumine.client.login.assert_called_with()
 
     def test_keep_alive_ka_error(self):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_client.betting_client.session_token = 1
-        mock_client.betting_client.session_expired = True
-        mock_client.betting_client.keep_alive.side_effect = BetfairError
-        worker.keep_alive(mock_context, mock_client)
-        mock_client.keep_alive.assert_called_with()
+        mock_flumine = mock.Mock()
+        mock_flumine.client.betting_client.session_token = 1
+        mock_flumine.client.betting_client.session_expired = True
+        mock_flumine.client.betting_client.keep_alive.side_effect = BetfairError
+        worker.keep_alive(mock_context, mock_flumine)
+        mock_flumine.client.keep_alive.assert_called_with()
 
     @mock.patch("flumine.worker.events")
     def test_poll_market_catalogue(self, mock_events):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_markets = mock.Mock()
-        mock_markets.markets = {"1.234": None, "5.678": None}
-        mock_handler_queue = mock.Mock()
+        mock_flumine = mock.Mock()
+        mock_flumine.markets.markets = {"1.234": None, "5.678": None}
 
-        worker.poll_market_catalogue(
-            mock_context, mock_client, mock_markets, mock_handler_queue
-        )
-        mock_client.betting_client.betting.list_market_catalogue.assert_called_with(
-            filter={"marketIds": list(mock_markets.markets.keys())},
+        worker.poll_market_catalogue(mock_context, mock_flumine)
+        mock_flumine.client.betting_client.betting.list_market_catalogue.assert_called_with(
+            filter={"marketIds": list(mock_flumine.markets.markets.keys())},
             market_projection=[
                 "COMPETITION",
                 "EVENT",
@@ -92,18 +96,19 @@ class WorkersTest(unittest.TestCase):
             ],
             max_results=100,
         )
-        mock_handler_queue.put.assert_called_with(mock_events.MarketCatalogueEvent())
+        mock_flumine.handler_queue.put.assert_called_with(
+            mock_events.MarketCatalogueEvent()
+        )
 
     @mock.patch("flumine.worker.events")
     def test_poll_account_balance(self, mock_events):
         mock_context = mock.Mock()
-        mock_client = mock.Mock()
-        mock_client.account_funds = {1: 2}
         mock_flumine = mock.Mock()
-        worker.poll_account_balance(mock_context, mock_flumine, mock_client)
-        mock_client.update_account_details.assert_called_with()
+        mock_flumine.client.account_funds = {1: 2}
+        worker.poll_account_balance(mock_context, mock_flumine)
+        mock_flumine.client.update_account_details.assert_called_with()
         mock_flumine.log_control.assert_called_with(
-            mock_events.BalanceEvent(mock_client.account_funds)
+            mock_events.BalanceEvent(mock_flumine.client.account_funds)
         )
 
     # @mock.patch("flumine.worker._get_cleared_market")