Django单元测试等待数据库

import time from django.db import connections from django.db.utils import OperationalError from django.core.management.base import BaseCommand class Command(BaseCommand): """Django command to pause execution until database is available""" def handle(self, *args, **options): """Handle the command""" self.stdout.write('Waiting for database...') db_conn = None while not db_conn: try: db_conn = connections['default'] except OperationalError: self.stdout.write('Database unavailable, waiting 1 second...') time.sleep(0.1) self.stdout.write(self.style.SUCCESS('Database available!'))

def test_wait_for_db_ready(self): """Test waiting for db when db is available""" with patch('django.db.utils.ConnectionHandler.__getitem__') as gi: gi.return_value = True call_command('wait_for_db') self.assertTrue(True)

def test_wait_for_db(self): """Test waiting for db""" attempt = 0 def getitem(alias): if attempt < 5: attempt += 1 raise OperationalError() else: return True with patch('django.db.utils.ConnectionHandler.__getitem__') as gi: gi.side_effect = getitem call_command('wait_for_db') self.assertGreaterEqual(attempt, 5)

实现这一点有几种方法。最简单的方法可能就是放弃 getitem() 嵌套函数并使用 OperationalError s.然后您可以用补丁验证尝试的次数。 gi 对象的 call_count . 例如:

def test_wait_for_db(self):
    """Test waiting for db"""

    with patch('django.db.utils.ConnectionHandler.__getitem__') as gi:
        gi.side_effect = [OperationalError] * 5 + [True]
        call_command('wait_for_db')
        self.assertGreaterEqual(gi.call_count, 5)  # Verify using the call_count

如果你想保留 Ge() 功能,那么我想你只需要 attempt 变量 nonlocal 因此可以在嵌套函数中看到:

def test_wait_for_db(self):
    """Test waiting for db"""
    attempt = 0

    def getitem(alias):
        nonlocal attempt  # Make the outer attempt variable visible
        if attempt < 5:
            attempt += 1
            raise OperationalError()
        else:
            return True

    with patch('django.db.utils.ConnectionHandler.__getitem__') as gi:
        gi.side_effect = getitem
        call_command('wait_for_db')
        self.assertGreaterEqual(attempt, 5)

第三,正如评论中所建议的,您可以创建一个具有 尝试 属性,并使用类的实例作为副作用:

def test_wait_for_db(self):
    """Test waiting for db"""

    class Getitem:
        def __init__(self):
            self.attempt = 0

        def __call__(self, item):
            if self.attempt < 5:
                self.attempt += 1
                raise OperationalError()
            else:
                return True

    with patch('django.db.utils.ConnectionHandler.__getitem__') as gi:
        getitem = Getitem()
        gi.side_effect = getitem
        call_command('wait_for_db')
        self.assertGreaterEqual(getitem.attempt, 5)  # Access the attempts from the instance