请求解析Django项目等待PostgreSQL就绪命令的测试代码逻辑

背景 #

我正在学习一门Django在线课程，课程中需要将项目与PostgreSQL数据库连接。为处理Django应用先于PostgreSQL启动的场景，编写了一个延迟应用启动直至数据库就绪的命令，代码如下：

from psycopg2 import OperationalError as Psycopg2OpError

from django.db.utils import OperationalError
from django.core.management.base import BaseCommand
import time  # 注：原代码遗漏time模块导入，此处补充

class Command(BaseCommand):

    def handle(self, *args, **options):
        """ Entrypoint for command """
        self.stdout.write('Waiting for database...')
        db_up = False
        while db_up is False:
            try:
                self.check(databases=['default'])
                db_up = True
            except(Psycopg2OpError, OperationalError):
                self.stdout.write('Database unavailable, waiting for 1 second...')
                time.sleep(1)

        self.stdout.write(self.style.SUCCESS('Database available!'))

测试代码逻辑解析 #

下面是用于验证该命令有效性的测试代码，核心通过Mock模拟技术避免真实连接数据库，仅测试命令的逻辑行为：

from unittest.mock import patch

from psycopg2 import OperationalError as Psycopg2Error

from django.core.management import call_command
from django.db.utils import OperationalError
from django.test import SimpleTestCase

@patch('core.management.commands.wait_for_db.Command.check')
class CommandTests(SimpleTestCase):
    """ Test commands """

    def test_wait_for_db_ready(self, patched_check):
        """ Test waiting for database if database is ready"""
        patched_check.return_value = True

        call_command('wait_for_db')

        patched_check.assert_called_once_with(databases=['default'])


    @patch('time.sleep')
    def test_wait_for_db_delay(self, patched_sleep, patched_check):
        """ Test waiting for database when getting OperationalError"""

        patched_check.side_effect = [Psycopg2Error] * 2 + \
            [OperationalError] * 3 + [True]

        call_command('wait_for_db')

        self.assertEqual(patched_check.call_count, 6)
        patched_check.assert_called_with(databases=['default'])

1. 测试类与全局Mock #

@patch('core.management.commands.wait_for_db.Command.check') 对命令中的self.check()方法做全局Mock，所有测试用例都不会真正执行数据库检查逻辑，而是用模拟的返回值/异常替代。

2. 第一个测试用例：数据库已就绪的情况 #

• test_wait_for_db_ready 验证数据库一开始就可用时的命令行为：
  • patched_check.return_value = True：模拟check()方法直接返回True，代表数据库状态正常
  • call_command('wait_for_db')：调用我们编写的wait_for_db管理命令
  • patched_check.assert_called_once_with(databases=['default'])：断言check()仅被调用1次，且传入参数为databases=['default']，确保命令正确指定了要检查的数据库实例

3. 第二个测试用例：数据库延迟就绪的情况 #

• test_wait_for_db_delay 验证数据库初始不可用时，命令反复重试直至就绪的逻辑：
  • 内层@patch('time.sleep')：Mocktime.sleep()方法，避免测试过程中真实等待，提升测试效率
  • patched_check.side_effect = [Psycopg2Error]*2 + [OperationalError]*3 + [True]：模拟check()的调用结果：前2次抛出PostgreSQL底层连接异常，接下来3次抛出Django层面的数据库不可用异常，第6次调用返回True（数据库就绪）
  • call_command('wait_for_db')：调用命令
  • self.assertEqual(patched_check.call_count, 6)：断言check()总共被调用6次（前5次失败后重试，第6次成功跳出循环）
  • patched_check.assert_called_with(databases=['default'])：断言每次调用check()都传入了正确的参数databases=['default']

内容的提问来源于stack exchange，提问作者ahmet yılmaz234