ATOCore/tests/test_engineering.py

"""Tests for the Engineering Knowledge Layer."""

from atocore.engineering.service import (
    ENTITY_TYPES,
    RELATIONSHIP_TYPES,
    create_entity,
    create_relationship,
    get_entities,
    get_entity,
    get_entity_with_context,
    get_relationships,
    init_engineering_schema,
)
from atocore.models.database import init_db

import pytest


def test_create_and_get_entity(tmp_data_dir):
    init_db()
    init_engineering_schema()
    e = create_entity(
        entity_type="component",
        name="Pivot Pin",
        project="p04-gigabit",
        description="Lateral support pivot pin for M1 assembly",
        properties={"material": "GF-PTFE", "diameter_mm": 12},
    )
    assert e.entity_type == "component"
    assert e.name == "Pivot Pin"
    assert e.properties["material"] == "GF-PTFE"

    fetched = get_entity(e.id)
    assert fetched is not None
    assert fetched.name == "Pivot Pin"


def test_create_relationship(tmp_data_dir):
    init_db()
    init_engineering_schema()
    subsystem = create_entity("subsystem", "Lateral Support", project="p04-gigabit")
    component = create_entity("component", "Pivot Pin", project="p04-gigabit")

    rel = create_relationship(
        source_entity_id=subsystem.id,
        target_entity_id=component.id,
        relationship_type="contains",
    )
    assert rel.relationship_type == "contains"

    rels = get_relationships(subsystem.id, direction="outgoing")
    assert len(rels) == 1
    assert rels[0].target_entity_id == component.id


def test_entity_with_context(tmp_data_dir):
    init_db()
    init_engineering_schema()
    subsystem = create_entity("subsystem", "Lateral Support", project="p04-gigabit")
    pin = create_entity("component", "Pivot Pin", project="p04-gigabit")
    pad = create_entity("component", "PTFE Pad", project="p04-gigabit")
    material = create_entity("material", "GF-PTFE", project="p04-gigabit",
                             description="Glass-filled PTFE for thermal stability")

    create_relationship(subsystem.id, pin.id, "contains")
    create_relationship(subsystem.id, pad.id, "contains")
    create_relationship(pad.id, material.id, "uses_material")

    ctx = get_entity_with_context(subsystem.id)
    assert ctx is not None
    assert len(ctx["relationships"]) == 2
    assert pin.id in ctx["related_entities"]
    assert pad.id in ctx["related_entities"]


def test_filter_entities_by_type_and_project(tmp_data_dir):
    init_db()
    init_engineering_schema()
    create_entity("component", "Pin A", project="p04-gigabit")
    create_entity("component", "Pin B", project="p04-gigabit")
    create_entity("material", "Steel", project="p04-gigabit")
    create_entity("component", "Actuator", project="p06-polisher")

    components = get_entities(entity_type="component", project="p04-gigabit")
    assert len(components) == 2

    all_p04 = get_entities(project="p04-gigabit")
    assert len(all_p04) == 3

    polisher = get_entities(project="p06-polisher")
    assert len(polisher) == 1


def test_invalid_entity_type_raises(tmp_data_dir):
    init_db()
    init_engineering_schema()
    with pytest.raises(ValueError, match="Invalid entity type"):
        create_entity("spaceship", "Enterprise")


def test_invalid_relationship_type_raises(tmp_data_dir):
    init_db()
    init_engineering_schema()
    a = create_entity("component", "A")
    b = create_entity("component", "B")
    with pytest.raises(ValueError, match="Invalid relationship type"):
        create_relationship(a.id, b.id, "loves")


def test_entity_name_search(tmp_data_dir):
    init_db()
    init_engineering_schema()
    create_entity("component", "Vertical Support Pad")
    create_entity("component", "Lateral Support Bracket")
    create_entity("component", "Reference Frame")

    results = get_entities(name_contains="Support")
    assert len(results) == 2