210 lines
6.8 KiB
Python
210 lines
6.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""EtlTimer 单元测试"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import time
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
from zoneinfo import ZoneInfo
|
||
|
||
import pytest
|
||
|
||
from utils.timer import EtlTimer, StepRecord, _fmt_ms
|
||
|
||
|
||
# ── _fmt_ms 格式化 ───────────────────────────────────────────
|
||
|
||
|
||
class TestFmtMs:
|
||
def test_sub_second(self):
|
||
assert _fmt_ms(123.4) == "123.4ms"
|
||
|
||
def test_seconds(self):
|
||
assert _fmt_ms(2500) == "2.50s"
|
||
|
||
def test_minutes(self):
|
||
result = _fmt_ms(90_000) # 90 秒
|
||
assert result.startswith("1m")
|
||
|
||
def test_hours(self):
|
||
result = _fmt_ms(3_700_000) # ~61.7 分钟
|
||
assert result.startswith("1h")
|
||
|
||
|
||
# ── StepRecord ───────────────────────────────────────────────
|
||
|
||
|
||
class TestStepRecord:
|
||
def test_elapsed_seconds(self):
|
||
rec = StepRecord(name="test", start_time=None, elapsed_ms=1500.0)
|
||
assert rec.elapsed_seconds == 1.5
|
||
|
||
def test_to_dict_without_end(self):
|
||
from datetime import datetime
|
||
|
||
now = datetime.now(ZoneInfo("Asia/Shanghai"))
|
||
rec = StepRecord(name="s1", start_time=now, elapsed_ms=100.0)
|
||
d = rec.to_dict()
|
||
assert d["name"] == "s1"
|
||
assert d["end_time"] is None
|
||
assert d["elapsed_ms"] == 100.0
|
||
assert d["children"] == []
|
||
|
||
def test_to_dict_with_children(self):
|
||
from datetime import datetime
|
||
|
||
tz = ZoneInfo("Asia/Shanghai")
|
||
now = datetime.now(tz)
|
||
parent = StepRecord(name="p", start_time=now, elapsed_ms=200.0)
|
||
child = StepRecord(name="c", start_time=now, end_time=now, elapsed_ms=50.0)
|
||
parent.children.append(child)
|
||
d = parent.to_dict()
|
||
assert len(d["children"]) == 1
|
||
assert d["children"][0]["name"] == "c"
|
||
|
||
|
||
# ── EtlTimer 核心流程 ────────────────────────────────────────
|
||
|
||
|
||
class TestEtlTimer:
|
||
def test_start_stop_step(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("STEP_A")
|
||
time.sleep(0.01) # 确保有可测量的耗时
|
||
rec = timer.stop_step("STEP_A")
|
||
|
||
assert rec.name == "STEP_A"
|
||
assert rec.end_time is not None
|
||
assert rec.elapsed_ms > 0
|
||
|
||
def test_sub_steps(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("PARENT")
|
||
timer.start_sub_step("PARENT", "child_1")
|
||
time.sleep(0.01)
|
||
timer.stop_sub_step("PARENT", "child_1")
|
||
timer.start_sub_step("PARENT", "child_2")
|
||
timer.stop_sub_step("PARENT", "child_2")
|
||
timer.stop_step("PARENT")
|
||
|
||
parent = timer.get_step("PARENT")
|
||
assert parent is not None
|
||
assert len(parent.children) == 2
|
||
assert parent.children[0].name == "child_1"
|
||
assert parent.children[0].elapsed_ms > 0
|
||
|
||
def test_stop_unknown_step_raises(self):
|
||
timer = EtlTimer()
|
||
with pytest.raises(KeyError, match="未找到步骤"):
|
||
timer.stop_step("NONEXISTENT")
|
||
|
||
def test_start_sub_step_unknown_parent_raises(self):
|
||
timer = EtlTimer()
|
||
with pytest.raises(KeyError, match="未找到父步骤"):
|
||
timer.start_sub_step("NONEXISTENT", "child")
|
||
|
||
def test_stop_sub_step_unknown_raises(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("P")
|
||
with pytest.raises(KeyError, match="未找到子步骤"):
|
||
timer.stop_sub_step("P", "no_such_child")
|
||
|
||
def test_multiple_steps(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
for name in ["ODS_LOAD", "DWD_LOAD", "DWS_AGG"]:
|
||
timer.start_step(name)
|
||
timer.stop_step(name)
|
||
|
||
assert len(timer.steps) == 3
|
||
assert timer.steps[0].name == "ODS_LOAD"
|
||
assert timer.steps[2].name == "DWS_AGG"
|
||
|
||
def test_to_dict(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("S1")
|
||
timer.stop_step("S1")
|
||
report = timer.finish(write_report=False)
|
||
|
||
d = timer.to_dict()
|
||
assert d["overall_start"] is not None
|
||
assert d["overall_end"] is not None
|
||
assert d["overall_elapsed_ms"] >= 0
|
||
assert len(d["steps"]) == 1
|
||
|
||
def test_overall_elapsed(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
time.sleep(0.02)
|
||
timer.finish(write_report=False)
|
||
assert timer.overall_elapsed_ms >= 15 # 至少 15ms（留余量）
|
||
|
||
def test_finish_returns_markdown(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("TEST_STEP")
|
||
timer.stop_step("TEST_STEP")
|
||
md = timer.finish(write_report=False)
|
||
|
||
assert "# ETL 执行计时报告" in md
|
||
assert "TEST_STEP" in md
|
||
assert "步骤汇总" in md
|
||
|
||
def test_markdown_contains_sub_steps(self):
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("MAIN")
|
||
timer.start_sub_step("MAIN", "sub_a")
|
||
timer.stop_sub_step("MAIN", "sub_a")
|
||
timer.stop_step("MAIN")
|
||
md = timer.finish(write_report=False)
|
||
|
||
assert "步骤详情" in md
|
||
assert "sub_a" in md
|
||
|
||
def test_write_report_requires_env(self):
|
||
"""ETL_REPORT_ROOT 未设置时应抛出 KeyError"""
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("X")
|
||
timer.stop_step("X")
|
||
|
||
with patch.dict(os.environ, {}, clear=True):
|
||
with pytest.raises(KeyError, match="ETL_REPORT_ROOT"):
|
||
timer.finish(write_report=True)
|
||
|
||
def test_write_report_creates_file(self, tmp_path: Path):
|
||
"""设置 ETL_REPORT_ROOT 后应生成 .md 文件"""
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("Y")
|
||
timer.stop_step("Y")
|
||
|
||
with patch.dict(os.environ, {"ETL_REPORT_ROOT": str(tmp_path)}):
|
||
timer.finish(write_report=True)
|
||
|
||
md_files = list(tmp_path.glob("etl_timing_*.md"))
|
||
assert len(md_files) == 1
|
||
content = md_files[0].read_text(encoding="utf-8")
|
||
assert "# ETL 执行计时报告" in content
|
||
assert "Y" in content
|
||
|
||
def test_elapsed_equals_end_minus_start(self):
|
||
"""Property 7 核心验证：耗时 ≈ 结束时间 - 开始时间"""
|
||
timer = EtlTimer()
|
||
timer.start()
|
||
timer.start_step("VERIFY")
|
||
time.sleep(0.05)
|
||
rec = timer.stop_step("VERIFY")
|
||
|
||
# 用 datetime 差值计算的毫秒数
|
||
dt_diff_ms = (rec.end_time - rec.start_time).total_seconds() * 1000
|
||
# perf_counter 计算的毫秒数
|
||
# 两者应在合理误差范围内（±50ms，考虑系统调度抖动）
|
||
assert abs(rec.elapsed_ms - dt_diff_ms) < 50
|