MajorMatch/tests/test_predict.py at main · eigenlambda123/MajorMatch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from dataclasses import dataclass

import api.predict as predict_api


@dataclass
class FakeBundle:
    model: object
    encoder: object
    feature_columns: list[str]


class FakeModel:
    def __init__(self, predicted_class: int, probabilities: list[float]):
        self.predicted_class = predicted_class
        self.probabilities = probabilities
        self.last_input = None

    def predict(self, rows):
        self.last_input = rows
        return [self.predicted_class]

    def predict_proba(self, rows):
        self.last_input = rows
        return [self.probabilities]


class FakeEncoder:
    def __init__(self, labels: list[str]):
        self.labels = labels

    def inverse_transform(self, values):
        return [self.labels[int(values[0])]]


def test_predict_track_uses_model_bundle_and_maps_major_to_track(monkeypatch):
    fake_model = FakeModel(0, [0.84, 0.16])
    fake_encoder = FakeEncoder(["B.Tech.-Computer Science and Engineering", "BVA- Bachelor of Visual Arts"])
    fake_bundle = FakeBundle(
        model=fake_model,
        encoder=fake_encoder,
        feature_columns=["Coding", "Mathematics", "Designing", "Computer Parts", "Researching", "Solving Puzzles"],
    )

    monkeypatch.setattr(predict_api, "_load_artifacts_from_disk", lambda: fake_bundle)
    predict_api._get_model_bundle.cache_clear()

    prediction = predict_api.predict_track({"coding": 9, "math": 4, "design": 2})

    assert prediction["label"] == "B.Tech.-Computer Science and Engineering"
    assert prediction["category"] == "Software Engineer"
    assert prediction["confidence"] == 0.84
    assert len(prediction["top_predictions"]) == 2
    assert prediction["top_predictions"][0]["label"] == "B.Tech.-Computer Science and Engineering"
    assert fake_model.last_input is not None
    assert fake_model.last_input.iloc[0]["Coding"] == 1


def test_predict_track_falls_back_when_model_loading_fails(monkeypatch):
    monkeypatch.setattr(predict_api, "_load_artifacts_from_disk", lambda: (_ for _ in ()).throw(RuntimeError("boom")))
    predict_api._get_model_bundle.cache_clear()

    prediction = predict_api.predict_track({"coding": 9, "math": 4, "design": 2})

    # With the structured numeric profile removed, the fallback now
    # inspects selected feature names (dict keys) and picks a label by
    # keyword matching. The provided keys include 'math', so the
    # fallback is expected to pick Data Scientist.
    assert prediction["label"] == "Data Scientist"
    assert prediction["confidence"] == 0.5
    assert prediction["top_predictions"] == [
        {
            "label": "Data Scientist",
            "confidence": 0.5,
            "category": "Data Scientist",
        }
    ]


def test_predict_track_accepts_direct_feature_selections(monkeypatch):
    fake_model = FakeModel(1, [0.21, 0.79])
    fake_encoder = FakeEncoder(["B.Tech.-Computer Science and Engineering", "BVA- Bachelor of Visual Arts"])
    fake_bundle = FakeBundle(
        model=fake_model,
        encoder=fake_encoder,
        feature_columns=["Coding", "Mathematics", "Designing", "Computer Parts"],
    )

    monkeypatch.setattr(predict_api, "_load_artifacts_from_disk", lambda: fake_bundle)
    predict_api._get_model_bundle.cache_clear()

    prediction = predict_api.predict_track(["Coding", "Mathematics"])

    assert prediction["label"] == "BVA- Bachelor of Visual Arts"
    assert prediction["source"] == "model"
    assert len(prediction["top_predictions"]) == 2
    assert prediction["top_predictions"][0]["label"] == "BVA- Bachelor of Visual Arts"
    assert fake_model.last_input is not None
    assert fake_model.last_input.iloc[0]["Coding"] == 1
    assert fake_model.last_input.iloc[0]["Mathematics"] == 1