lemdb/tests/test_training_dataset.py

import sqlite3

import pytest

from lmpm.training.dataset import (
    TrainingDataError,
    load_sqlite_training_dataset,
)
from lmpm.training.workflow import train_sqlite_database

MATERIAL_SCHEMA = """
CREATE TABLE material_records (
    id INTEGER PRIMARY KEY,
    entry_batch_id INTEGER NOT NULL,
    experiment_id VARCHAR(100) NOT NULL UNIQUE,
    test_date DATE NOT NULL,
    material_category VARCHAR(100) NOT NULL,
    material_name VARCHAR(200) NOT NULL,
    apparent_density FLOAT,
    uv_absorption_355nm FLOAT,
    material_thickness FLOAT NOT NULL,
    moisture_content FLOAT,
    thermal_conductivity FLOAT,
    initial_decomposition_temp FLOAT,
    melting_vaporization_temp FLOAT,
    specific_heat_capacity FLOAT,
    carbon_residue_rate FLOAT,
    surface_roughness FLOAT,
    hardness FLOAT,
    actual_output_power FLOAT NOT NULL,
    display_current FLOAT NOT NULL,
    scanning_speed FLOAT NOT NULL,
    pulse_frequency FLOAT NOT NULL,
    pulse_width FLOAT NOT NULL,
    defocus_amount FLOAT NOT NULL,
    scan_line_spacing FLOAT NOT NULL,
    filling_method VARCHAR(100) NOT NULL,
    processing_size VARCHAR(100) NOT NULL,
    is_cut_through BOOLEAN NOT NULL,
    carbonized_edge_width FLOAT NOT NULL,
    etching_depth FLOAT NOT NULL,
    is_fire_smolder BOOLEAN NOT NULL,
    pattern_clarity_score FLOAT NOT NULL,
    presentation_balance_score FLOAT NOT NULL,
    finished_image_filename VARCHAR(255) NOT NULL,
    remarks TEXT NOT NULL,
    created_at DATETIME NOT NULL
);
"""


def make_record(number: int) -> tuple:
    return (
        number,
        1,
        f"LAS-2026-{number:04d}",
        "2026-09-11",
        "polymer",
        f"Acrylic {number}",
        1.18 if number == 1 else None,
        0.92,
        3.0,
        None,
        0.19,
        None,
        None,
        1.5,
        None,
        1.2,
        None,
        8.5 + number,
        2.0,
        80.0,
        20.0,
        5.0,
        0.0,
        0.1,
        "双向填充",
        "20x20",
        number == 1,
        0.2,
        30.0,
        False,
        9.0,
        8.0,
        f"image-{number}.jpg",
        "stable",
        "2026-09-11T08:00:00+00:00",
    )


def create_database(tmp_path, *, with_table: bool = True) -> object:
    database_path = tmp_path / f"platform-{int(with_table)}.db"
    connection = sqlite3.connect(database_path)
    if with_table:
        connection.execute(MATERIAL_SCHEMA)
        connection.executemany(
            """
            INSERT INTO material_records VALUES (
                ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
                ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
            )
            """,
            [make_record(1), make_record(2)],
        )
    connection.commit()
    connection.close()
    return database_path


def test_sqlite_dataset_maps_features_and_targets(tmp_path):
    dataset = load_sqlite_training_dataset(create_database(tmp_path))

    assert dataset.experiment_ids == ["LAS-2026-0001", "LAS-2026-0002"]
    assert dataset.feature_names[:5] == [
        "apparent_density",
        "uv_absorption_355nm",
        "material_thickness",
        "moisture_content",
        "thermal_conductivity",
    ]
    assert dataset.feature_names[-1] == "processing_size"
    assert dataset.target_names == [
        "is_cut_through",
        "carbonized_edge_width",
        "etching_depth",
        "is_fire_smolder",
        "pattern_clarity_score",
        "presentation_balance_score",
    ]
    assert dataset.records[0].metadata == {
        "experiment_id": "LAS-2026-0001",
        "test_date": "2026-09-11",
        "material_category": "polymer",
        "material_name": "Acrylic 1",
    }
    assert dataset.records[0].features["actual_output_power"] == pytest.approx(9.5)
    assert dataset.records[0].targets["is_cut_through"] is True
    assert dataset.records[1].targets["is_cut_through"] is False


def test_missing_database_raises_expected_error(tmp_path):
    missing_path = tmp_path / "missing.db"

    with pytest.raises(FileNotFoundError):
        load_sqlite_training_dataset(missing_path)


def test_database_without_material_table_raises_expected_error(tmp_path):
    database_path = create_database(tmp_path, with_table=False)

    with pytest.raises(TrainingDataError, match="material_records"):
        load_sqlite_training_dataset(database_path)


def test_sqlite_workflow_trains_and_saves_every_target_model(tmp_path):
    database_path = create_database(tmp_path)
    artifact_path = tmp_path / "lmpm.joblib"

    bundle = train_sqlite_database(database_path, artifact_path)

    assert artifact_path.exists()
    assert set(bundle.models) == {
        "is_cut_through",
        "carbonized_edge_width",
        "etching_depth",
        "is_fire_smolder",
        "pattern_clarity_score",
        "presentation_balance_score",
    }