Question

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

Answer 1

有几种方法可以实现这一目标。最简单的方法可能就是放弃getitem()嵌套函数，并使用一系列OperationalError来设置副作用。然后，您可以使用修补后的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

如果您希望保留getitem()函数，那么我认为您只需要制作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)

第三，根据注释的建议，您可以创建一个具有attempt属性的类，并使用该类的实例作为副作用：

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

Answer 2

您可以使用以下代码以更有效的方式实现同样的目标。

from unittest.mock import patch

from django.core.management import call_command
from django.db.utils import OperationalError
#gives error when db isn't available
from django.test import TestCase


class CommandTests(TestCase):

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

    @patch('time.sleep', return_value=True)
    def test_wait_for_db(self, ts):
        """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.assertEqual(gi.call_count, 6)

Django单元测试等待数据库

2 个答案: