Модульный тест Django ждет базу данных

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 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

Если вы хотите сохранить функцию 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