From 581e7a80f75191b8eab3356bd0b80daacd4db42a Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Thu, 16 Sep 2021 19:16:40 +0200
Subject: [PATCH] add uwsgidecorators module (#57019)

---
 hobo/multitenant/uwsgidecorators.py       | 53 +++++++++++++++++++++++
 tests_multitenant/test_uwsgidecorators.py | 41 ++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 hobo/multitenant/uwsgidecorators.py
 create mode 100644 tests_multitenant/test_uwsgidecorators.py

diff --git a/hobo/multitenant/uwsgidecorators.py b/hobo/multitenant/uwsgidecorators.py
new file mode 100644
index 0000000..5b0727a
--- /dev/null
+++ b/hobo/multitenant/uwsgidecorators.py
@@ -0,0 +1,53 @@
+# Copyright (C) 2021 Entr'ouvert
+
+import logging
+import pickle
+
+try:
+    import uwsgi
+except ImportError:
+    uwsgi = None
+
+logger = logging.getLogger(__name__)
+
+spooler_registry = {}
+
+
+def spool(func):
+    name = '%s.%s' % (func.__module__, func.__name__)
+    spooler_registry[name] = func
+
+    def spool_function(*args, **kwargs):
+        uwsgi.spool(name=name.encode(), body=pickle.dumps({'args': args, 'kwargs': kwargs}))
+
+    func.spool = spool_function
+    return func
+
+
+if uwsgi:
+
+    def spooler_function(env):
+        try:
+            try:
+                name = env.get('name').decode()
+                body = env.get('body')
+            except Exception:
+                logger.error('spooler: no name or body found: env.keys()=%s', env.keys())
+                return uwsgi.SPOOL_OK
+            try:
+                params = pickle.loads(body)
+                args = params['args']
+                kwargs = params['kwargs']
+            except Exception:
+                logger.exception('spooler: depickling of body failed')
+                return uwsgi.SPOOL_OK
+            try:
+                function = spooler_registry[name]
+            except KeyError:
+                logger.error('spooler: no function named "%s"', name)
+            function(*args, **kwargs)
+        except Exception:
+            logger.exception('spooler: function "%s" raied')
+        return uwsgi.SPOOL_OK
+
+    uwsgi.spooler = spooler_function
diff --git a/tests_multitenant/test_uwsgidecorators.py b/tests_multitenant/test_uwsgidecorators.py
new file mode 100644
index 0000000..84411fb
--- /dev/null
+++ b/tests_multitenant/test_uwsgidecorators.py
@@ -0,0 +1,41 @@
+import importlib
+import pickle
+
+import mock
+import pytest
+
+import hobo.multitenant.uwsgidecorators
+
+
+@pytest.fixture
+def uwsgi():
+    import sys
+
+    uwsgi = mock.Mock()
+    uwsgi.SPOOL_OK = -2
+    sys.modules['uwsgi'] = uwsgi
+    importlib.reload(hobo.multitenant.uwsgidecorators)
+    yield uwsgi
+    del sys.modules['uwsgi']
+
+
+def test_basic():
+    @hobo.multitenant.uwsgidecorators.spool
+    def function(a, b):
+        pass
+
+    function(1, 2)
+    with pytest.raises(AttributeError):
+        function.spool(1, 2)
+
+
+def test_mocked_uwsgi(uwsgi):
+    @hobo.multitenant.uwsgidecorators.spool
+    def function(a, b):
+        pass
+
+    function(1, 2)
+    function.spool(1, 2)
+    assert set(uwsgi.spool.call_args[1].keys()) == {'body', 'name'}
+    assert pickle.loads(uwsgi.spool.call_args[1]['body']) == {'args': (1, 2), 'kwargs': {}}
+    assert uwsgi.spool.call_args[1]['name'] == b'test_uwsgidecorators.function'
-- 
2.33.0