Produits Entr'ouvert » Hobo

# 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

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'