Python使用pytest高效编写和管理单元测试的完整指南

import unittest
from my_module import add

class TestAdd(unittest.TestCase):
    def test_add_positive_numbers(self):
        self.assertEqual(add(1, 2), 3)

from my_module import add

def test_add_positive_numbers():
    assert add(1, 2) == 3

import pytest
from my_module import is_prime

@pytest.mark.parametrize("number, expected", [
    (2, True),
    (3, True),
    (4, False),
    (10, False),
    (17, True)
])
def test_is_prime(number, expected):
    assert is_prime(number) == expected

import pytest

@pytest.fixture
def db_connection():
    # Setup: 建立连接
    print("\n连接数据库...")
    connection = {"status": "connected"}
    yield connection
    # Teardown: 关闭连接
    print("\n断开数据库...")

def test_db_query(db_connection):
    # 测试函数直接使用 fixture 返回的对象
    assert db_connection["status"] == "connected"

import pytest
import psycopg2

# 数据库连接配置 (建议从环境变量读取)
DB_CONFIG = {
    "dbname": "test_db",
    "user": "postgres",
    "password": "password",
    "host": "localhost"
}

@pytest.fixture(scope="function")
def db_cursor():
    """
    创建数据库连接和游标，并在测试结束后回滚。
    scope="function" 表示每个测试函数都会运行一次此 fixture。
    """
    conn = psycopg2.connect(**DB_CONFIG)
    cursor = conn.cursor()
    
    # 开始事务
    cursor.execute("BEGIN")
    
    yield cursor
    
    # 测试结束，回滚事务
    cursor.execute("ROLLBACK")
    cursor.close()
    conn.close()

def test_insert_user(db_cursor):
    """
    测试插入用户逻辑。
    即使插入成功，由于事务回滚，数据库中不会留下这条记录。
    """
    # 假设有一个 users 表
    db_cursor.execute("INSERT INTO users (name, email) VALUES (%s, %s)", ("Alice", "alice@example.com"))
    db_cursor.execute("SELECT count(*) FROM users WHERE name = %s", ("Alice",))
    count = db_cursor.fetchone()[0]
    
    assert count == 1

def test_query_user(db_cursor):
    """
    验证上一个测试没有影响当前环境。
    """
    db_cursor.execute("SELECT count(*) FROM users WHERE name = %s", ("Alice",))
    count = db_cursor.fetchone()[0]
    
    # 因为上一个测试回滚了，这里应该查不到 Alice
    assert count == 0

import socket
import pytest
from my_module import start_udp_listener

def test_udp_listener_logic(monkeypatch):
    # 模拟接收到的数据
    mock_data = b"Hello UDP"
    
    # 用来存储发送回客户端的数据
    sent_data = []
    
    class MockSocket:
        def __init__(self, family, type_):
            pass
        def bind(self, address):
            pass
        def recvfrom(self, bufsize):
            # 模拟接收到数据和发送方地址
            return (mock_data, ("127.0.0.1", 8080))
        def sendto(self, data, addr):
            sent_data.append(data)
        def close(self):
            pass

   # 使用 monkeypatch 替换 socket.socket
    monkeypatch.setattr(socket, "socket", MockSocket)

    # 执行被测代码 (注意：这里需要确保代码能被调用并阻塞，或者在单独线程运行)
    # 实际测试中，通常会将逻辑提取出来单独测试，或者使用超时机制
    # 这里仅作演示：假设 start_udp_listener 内部逻辑处理了 recvfrom 并调用 sendto
    
    # 模拟监听器的一次循环
    sock = MockSocket(socket.AF_INET, socket.SOCK_DGRAM)
    data, addr = sock.recvfrom(1024)
    # 假设处理逻辑是回显大写
    processed_data = data.upper()
    sock.sendto(processed_data, addr)
    
    assert sent_data[0] == b"HELLO UDP"