from pathlib import Path

from app.modules.weights.service import classify_weight


def test_classify_weight():
    item = classify_weight(Path("Seg_All_In_One_YoloModel/yolov8n-seg.pt"))
    assert item["family"] == "yolo"
    assert item["role"] == "weight"

    best = classify_weight(Path("Seg_Predict_YoloModel/YOLOv9e-seg/weights/best.pt"))
    assert best["family"] == "yolo"
    assert best["role"] == "trained_best"