tests/test_model.py

# unit test for the model
import joblib
import os 
import numpy as np

def test_model_loading():
    model_path = os.path.join("model","model.joblib")
    model = joblib.load(model_path)
    assert model is not None
    
def test_model_prediction():
    model_path = os.path.join("model","model.joblib")
    model = joblib.load(model_path)
    sample_test_data = np.array([[5.1], [3.5], [1.4], [0.2]])
    prediction = model.predict(sample_test_data)
    assert prediction[0] in [0,1,2]