332 lines
12 KiB
Python
332 lines
12 KiB
Python
"""
|
|
Tests for DatasetBuilder service.
|
|
|
|
TDD: Write tests first, then implement dataset_builder.py.
|
|
"""
|
|
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
|
|
from inference.data.admin_models import (
|
|
AdminAnnotation,
|
|
AdminDocument,
|
|
TrainingDataset,
|
|
FIELD_CLASSES,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def tmp_admin_images(tmp_path):
|
|
"""Create mock admin images directory with sample images."""
|
|
doc_ids = [uuid4() for _ in range(5)]
|
|
for doc_id in doc_ids:
|
|
doc_dir = tmp_path / "admin_images" / str(doc_id)
|
|
doc_dir.mkdir(parents=True)
|
|
# Create 2 pages per doc
|
|
for page in range(1, 3):
|
|
img_path = doc_dir / f"page_{page}.png"
|
|
img_path.write_bytes(b"fake-png-data")
|
|
return tmp_path, doc_ids
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_admin_db():
|
|
"""Mock AdminDB with dataset and document methods."""
|
|
db = MagicMock()
|
|
db.create_dataset.return_value = TrainingDataset(
|
|
dataset_id=uuid4(),
|
|
name="test-dataset",
|
|
status="building",
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
)
|
|
return db
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_documents(tmp_admin_images):
|
|
"""Create sample AdminDocument objects."""
|
|
tmp_path, doc_ids = tmp_admin_images
|
|
docs = []
|
|
for doc_id in doc_ids:
|
|
doc = MagicMock(spec=AdminDocument)
|
|
doc.document_id = doc_id
|
|
doc.filename = f"{doc_id}.pdf"
|
|
doc.page_count = 2
|
|
doc.file_path = str(tmp_path / "admin_images" / str(doc_id))
|
|
docs.append(doc)
|
|
return docs
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_annotations(sample_documents):
|
|
"""Create sample annotations for each document page."""
|
|
annotations = {}
|
|
for doc in sample_documents:
|
|
doc_anns = []
|
|
for page in range(1, 3):
|
|
ann = MagicMock(spec=AdminAnnotation)
|
|
ann.document_id = doc.document_id
|
|
ann.page_number = page
|
|
ann.class_id = 0
|
|
ann.class_name = "invoice_number"
|
|
ann.x_center = 0.5
|
|
ann.y_center = 0.3
|
|
ann.width = 0.2
|
|
ann.height = 0.05
|
|
doc_anns.append(ann)
|
|
annotations[str(doc.document_id)] = doc_anns
|
|
return annotations
|
|
|
|
|
|
class TestDatasetBuilder:
|
|
"""Tests for DatasetBuilder."""
|
|
|
|
def test_build_creates_directory_structure(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""Dataset builder should create images/ and labels/ with train/val/test subdirs."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
dataset_dir = tmp_path / "datasets" / "test"
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
|
|
# Mock DB calls
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
result_dir = tmp_path / "datasets" / str(dataset.dataset_id)
|
|
for split in ["train", "val", "test"]:
|
|
assert (result_dir / "images" / split).exists()
|
|
assert (result_dir / "labels" / split).exists()
|
|
|
|
def test_build_copies_images(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""Images should be copied from admin_images to dataset folder."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
result = builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
# Check total images copied
|
|
result_dir = tmp_path / "datasets" / str(dataset.dataset_id)
|
|
total_images = sum(
|
|
len(list((result_dir / "images" / split).glob("*.png")))
|
|
for split in ["train", "val", "test"]
|
|
)
|
|
assert total_images == 10 # 5 docs * 2 pages
|
|
|
|
def test_build_generates_yolo_labels(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""YOLO label files should be generated with correct format."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
result_dir = tmp_path / "datasets" / str(dataset.dataset_id)
|
|
total_labels = sum(
|
|
len(list((result_dir / "labels" / split).glob("*.txt")))
|
|
for split in ["train", "val", "test"]
|
|
)
|
|
assert total_labels == 10 # 5 docs * 2 pages
|
|
|
|
# Check label format: "class_id x_center y_center width height"
|
|
label_files = list((result_dir / "labels").rglob("*.txt"))
|
|
content = label_files[0].read_text().strip()
|
|
parts = content.split()
|
|
assert len(parts) == 5
|
|
assert int(parts[0]) == 0 # class_id
|
|
assert 0 <= float(parts[1]) <= 1 # x_center
|
|
assert 0 <= float(parts[2]) <= 1 # y_center
|
|
|
|
def test_build_generates_data_yaml(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""data.yaml should be generated with correct field classes."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
yaml_path = tmp_path / "datasets" / str(dataset.dataset_id) / "data.yaml"
|
|
assert yaml_path.exists()
|
|
content = yaml_path.read_text()
|
|
assert "train:" in content
|
|
assert "val:" in content
|
|
assert "nc:" in content
|
|
assert "invoice_number" in content
|
|
|
|
def test_build_splits_documents_correctly(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""Documents should be split into train/val/test according to ratios."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
# Verify add_dataset_documents was called with correct splits
|
|
call_args = mock_admin_db.add_dataset_documents.call_args
|
|
docs_added = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
|
|
splits = [d["split"] for d in docs_added]
|
|
assert "train" in splits
|
|
# With 5 docs, 80/10/10 -> 4 train, 0-1 val, 0-1 test
|
|
train_count = splits.count("train")
|
|
assert train_count >= 3 # At least 3 of 5 should be train
|
|
|
|
def test_build_updates_status_to_ready(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""After successful build, dataset status should be updated to 'ready'."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
mock_admin_db.update_dataset_status.assert_called_once()
|
|
call_kwargs = mock_admin_db.update_dataset_status.call_args[1]
|
|
assert call_kwargs["status"] == "ready"
|
|
assert call_kwargs["total_documents"] == 5
|
|
assert call_kwargs["total_images"] == 10
|
|
|
|
def test_build_sets_failed_on_error(
|
|
self, tmp_path, mock_admin_db
|
|
):
|
|
"""If build fails, dataset status should be set to 'failed'."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = [] # No docs found
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
with pytest.raises(ValueError):
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=["nonexistent-id"],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
|
|
mock_admin_db.update_dataset_status.assert_called_once()
|
|
call_kwargs = mock_admin_db.update_dataset_status.call_args[1]
|
|
assert call_kwargs["status"] == "failed"
|
|
|
|
def test_build_with_seed_produces_deterministic_splits(
|
|
self, tmp_path, mock_admin_db, sample_documents, sample_annotations
|
|
):
|
|
"""Same seed should produce same splits."""
|
|
from inference.web.services.dataset_builder import DatasetBuilder
|
|
|
|
results = []
|
|
for _ in range(2):
|
|
builder = DatasetBuilder(db=mock_admin_db, base_dir=tmp_path / "datasets")
|
|
mock_admin_db.get_documents_by_ids.return_value = sample_documents
|
|
mock_admin_db.get_annotations_for_document.side_effect = lambda doc_id: (
|
|
sample_annotations.get(str(doc_id), [])
|
|
)
|
|
mock_admin_db.add_dataset_documents.reset_mock()
|
|
mock_admin_db.update_dataset_status.reset_mock()
|
|
|
|
dataset = mock_admin_db.create_dataset.return_value
|
|
builder.build_dataset(
|
|
dataset_id=str(dataset.dataset_id),
|
|
document_ids=[str(d.document_id) for d in sample_documents],
|
|
train_ratio=0.8,
|
|
val_ratio=0.1,
|
|
seed=42,
|
|
admin_images_dir=tmp_path / "admin_images",
|
|
)
|
|
call_args = mock_admin_db.add_dataset_documents.call_args
|
|
docs = call_args[1]["documents"] if "documents" in call_args[1] else call_args[0][1]
|
|
results.append([(d["document_id"], d["split"]) for d in docs])
|
|
|
|
assert results[0] == results[1]
|