Files
Data_engineer_test_task/tests/test_pipeline.py
T

131 lines
4.6 KiB
Python

import csv
import json
import os
from pathlib import Path
import shutil
import sqlite3
import subprocess
import sys
from contextlib import closing
import pytest
from src.extract import FIELDS
ROOT = Path(__file__).resolve().parents[1]
ROWS = [
(8, "2024-01-01", "Книги", 2, "10,25"),
(3, "01.01.2024", " книги ", 1, "0.50"),
(5, "2024-01-01", "Дом", 3, "1"),
(2, "2024-01-02", "Книги", 1, "5"),
(1, "2024-01-02", "Дом", 0, "2"),
]
def write_input(project, rows):
with (project / "data/raw_orders.csv").open("w", encoding="utf-8", newline="") as out:
writer = csv.writer(out)
writer.writerow(FIELDS)
writer.writerows(rows)
@pytest.fixture
def project(tmp_path):
folder = tmp_path / "project"
(folder / "data").mkdir(parents=True)
shutil.copy(ROOT / "main.py", folder)
shutil.copytree(ROOT / "sql", folder / "sql")
shutil.copytree(ROOT / "src", folder / "src", ignore=shutil.ignore_patterns("__pycache__"))
write_input(folder, ROWS)
return folder
def run(project, *, success=True):
result = subprocess.run(
[sys.executable, str(project / "main.py")], cwd=project.parent,
env={**os.environ, "PYTHONIOENCODING": "utf-8"}, capture_output=True, encoding="utf-8",
)
if success:
assert result.returncode == 0, result.stderr
else:
assert result.returncode != 0
assert not result.stdout
return result
def snapshot(project):
with closing(sqlite3.connect(project / "outputs/orders.db")) as db:
orders = db.execute("SELECT * FROM orders ORDER BY order_id").fetchall()
daily = db.execute("SELECT * FROM daily_category_revenue ORDER BY order_date, category").fetchall()
return orders, daily
def test_full_run_and_repeat(project):
run(project)
orders, daily = snapshot(project)
assert len(orders) == 4
assert daily == [
("2024-01-01", "Дом", 1, 300), ("2024-01-01", "Книги", 2, 2100),
("2024-01-02", "Книги", 1, 500),
]
for day, category, count, revenue in daily:
amounts = [row[5] for row in orders if (row[1], row[2]) == (day, category)]
assert (len(amounts), sum(amounts)) == (count, revenue)
outputs = project / "outputs"
with (outputs / "clean_orders.csv").open(encoding="utf-8-sig", newline="") as source:
clean = list(csv.DictReader(source))
assert sorted((int(r["order_id"]), r["order_date"], r["category"], int(r["quantity"]),
int(r["unit_price_kopecks"]), int(r["revenue_kopecks"])) for r in clean) == orders
with (outputs / "rejected_rows.csv").open(encoding="utf-8-sig", newline="") as source:
assert [r["record_number"] for r in csv.DictReader(source)] == ["5"]
report = json.loads((outputs / "quality_report.json").read_text(encoding="utf-8"))
assert (report["source_count"], report["accepted_count"], report["rejected_count"]) == (5, 4, 1)
run(project)
assert snapshot(project) == (orders, daily)
def test_aggregate_overflow_rolls_back_both_tables(project):
run(project)
before = snapshot(project)
write_input(project, [
(99, "2025-01-01", "Дом", 1, "92233720368547758.07"),
(100, "2025-01-01", "Дом", 1, "0.01"),
])
assert "integer overflow" in run(project, success=False).stderr
assert snapshot(project) == before
report = json.loads((project / "outputs/quality_report.json").read_text(encoding="utf-8"))
assert report["source_count"] == len(ROWS)
def test_damaged_input_preserves_database(project):
run(project)
before = snapshot(project)
with (project / "data/raw_orders.csv").open("a", encoding="utf-8") as out:
out.write('123,2024-01-20,"незакрытое поле')
assert "ошибка структуры CSV" in run(project, success=False).stderr
assert snapshot(project) == before
def test_no_accepted_orders_clears_both_tables(project):
run(project)
write_input(project, [ROWS[-1]])
run(project)
assert snapshot(project) == ([], [])
report = json.loads((project / "outputs/quality_report.json").read_text(encoding="utf-8"))
assert report["rejected_count"] == 1
def test_export_error_after_commit(project):
run(project)
write_input(project, [ROWS[2]])
blocked = project / "outputs/clean_orders.csv"
blocked.unlink()
blocked.mkdir()
result = run(project, success=False)
assert "База уже обновлена" in result.stderr
assert "clean_orders.csv" in result.stderr
assert snapshot(project) == (
[(5, "2024-01-01", "Дом", 3, 100, 300)], [("2024-01-01", "Дом", 1, 300)]
)