import io
import joblib
from classifier.model import build_pipeline, predict


def test_train_and_predict_toy():
    pipe = build_pipeline()
    X = [
        "double extortion data leak",
        "single extortion encryption only",
    ]
    y = ["double-extortion", "single-extortion"]
    pipe.fit(X, y)
    buf = io.BytesIO()
    joblib.dump(pipe, buf)
    labels, conf, weights = predict(["double extortion leak"], buf.getvalue())
    assert labels[0] in {"double-extortion", "single-extortion"}
    assert 0.0 <= conf[0] <= 1.0
    assert len(weights[0]) > 0


