SMANAGE/gvsds/tests/test_celery_utils.py

"""模块：celery_utils.py 单元测试

覆盖 RollbackOnFailure / RetryOnFailure / LogTask / SafeDecimal
四个装饰器，无 Celery 环境下验证同步执行与异常处理。
"""
from decimal import Decimal
from unittest import mock

from django.test import SimpleTestCase

from gvsdsdk import celery_utils


class RollbackOnFailureTestCase(SimpleTestCase):
    """RollbackOnFailure 事务回滚装饰器。"""

    def testNormalReturn_PassesThrough(self):
        @celery_utils.RollbackOnFailure
        def Task():
            return 'ok'

        self.assertEqual(Task(), 'ok')

    def testException_MarksRollbackAndReRaises(self):
        @celery_utils.RollbackOnFailure
        def Task():
            raise ValueError('boom')

        with mock.patch('django.db.transaction.set_rollback') as mock_set_rollback:
            with self.assertRaises(ValueError):
                Task()
            mock_set_rollback.assert_called_once_with(True)

    def testException_TransactionSetRollbackFails_SilentlyReRaises(self):
        @celery_utils.RollbackOnFailure
        def Task():
            raise ValueError('boom')

        with mock.patch('django.db.transaction.set_rollback', side_effect=RuntimeError):
            with self.assertRaises(ValueError):
                Task()  # 不应抛 RuntimeError

    def testPreservesFuncMetadata(self):
        @celery_utils.RollbackOnFailure
        def ImportantTask():
            """Important."""

        self.assertEqual(ImportantTask.__name__, 'ImportantTask')
        self.assertEqual(ImportantTask.__doc__, 'Important.')


class RetryOnFailureTestCase(SimpleTestCase):
    """RetryOnFailure 重试装饰器。"""

    def testSuccessOnFirstAttempt_NoRetry(self):
        CallCount = {'n': 0}

        @celery_utils.RetryOnFailure(max_retries=3, delay=0)
        def Task():
            CallCount['n'] += 1
            return 'ok'

        self.assertEqual(Task(), 'ok')
        self.assertEqual(CallCount['n'], 1)

    def testSuccessAfterRetries(self):
        CallCount = {'n': 0}

        @celery_utils.RetryOnFailure(max_retries=3, delay=0)
        def Task():
            CallCount['n'] += 1
            if CallCount['n'] < 3:
                raise ValueError('retry')
            return 'ok'

        self.assertEqual(Task(), 'ok')
        self.assertEqual(CallCount['n'], 3)

    def testAllRetriesExhausted_RaisesLastException(self):
        CallCount = {'n': 0}

        @celery_utils.RetryOnFailure(max_retries=2, delay=0)
        def Task():
            CallCount['n'] += 1
            raise ValueError('always fail')

        with self.assertRaises(ValueError):
            Task()
        # 1 次初始 + 2 次重试 = 3 次
        self.assertEqual(CallCount['n'], 3)

    def testOnlyRetriesOnSpecifiedExceptions(self):
        CallCount = {'n': 0}

        @celery_utils.RetryOnFailure(max_retries=3, delay=0, exceptions=(ValueError,))
        def Task():
            CallCount['n'] += 1
            raise TypeError('not retried')

        with self.assertRaises(TypeError):
            Task()
        # TypeError 不在重试列表 → 立即抛出，只调用 1 次
        self.assertEqual(CallCount['n'], 1)

    def testZeroMaxRetries_NoRetry(self):
        CallCount = {'n': 0}

        @celery_utils.RetryOnFailure(max_retries=0, delay=0)
        def Task():
            CallCount['n'] += 1
            raise ValueError('fail')

        with self.assertRaises(ValueError):
            Task()
        self.assertEqual(CallCount['n'], 1)


class LogTaskTestCase(SimpleTestCase):
    """LogTask 日志装饰器。"""

    def testSuccess_ReturnsResult(self):
        @celery_utils.LogTask
        def Task():
            return 42

        self.assertEqual(Task(), 42)

    def testException_ReRaises(self):
        @celery_utils.LogTask
        def Task():
            raise RuntimeError('boom')

        with self.assertRaises(RuntimeError):
            Task()

    def testPreservesFuncMetadata(self):
        @celery_utils.LogTask
        def NamedTask():
            """Named."""

        self.assertEqual(NamedTask.__name__, 'NamedTask')


class SafeDecimalTestCase(SimpleTestCase):
    """SafeDecimal 安全运算装饰器。"""

    def testNormalReturn_PassesThrough(self):
        @celery_utils.SafeDecimal(default=Decimal('-1'))
        def Calc():
            return Decimal('100')

        self.assertEqual(Calc(), Decimal('100'))

    def testArithmeticError_ReturnsDefault(self):
        @celery_utils.SafeDecimal(default=Decimal('-1'))
        def Calc():
            raise ArithmeticError('div by zero')

        self.assertEqual(Calc(), Decimal('-1'))

    def testTypeError_ReturnsDefault(self):
        @celery_utils.SafeDecimal(default=Decimal('0'))
        def Calc():
            raise TypeError('bad type')

        self.assertEqual(Calc(), Decimal('0'))

    def testValueError_ReturnsDefault(self):
        @celery_utils.SafeDecimal()
        def Calc():
            raise ValueError('bad value')

        self.assertEqual(Calc(), Decimal('0'))

    def testNonListedException_ReRaises(self):
        @celery_utils.SafeDecimal()
        def Calc():
            raise RuntimeError('not caught')

        with self.assertRaises(RuntimeError):
            Calc()


class IsCeleryTaskTestCase(SimpleTestCase):
    """_IsCeleryTask 判断逻辑。"""

    def testPlainFunc_ReturnsFalseWhenNoCelery(self):
        def PlainFunc():
            pass

        # 无 Celery 时返回 False
        with mock.patch.object(celery_utils, '_HAS_CELERY', False):
            self.assertFalse(celery_utils._IsCeleryTask(PlainFunc))

    def testFuncWithApplyAsync_ReturnsTrue(self):
        class FakeTask:
            def apply_async(self):
                pass

        with mock.patch.object(celery_utils, '_HAS_CELERY', True):
            self.assertTrue(celery_utils._IsCeleryTask(FakeTask()))