Django单元测试等待数据库

我有一个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)

是否有一种方法可以测试命令在返回之前是否等待数据库可用?

到目前为止,我已经尝试了以下方法,但是它不起作用 attempt 无法从外部访问 getitem .

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)