public-engines/product-classifier-engine/tests/prediction/test_predictor.py

#!/usr/bin/env python
# coding=utf-8

try:
    import mock

except ImportError:
    import unittest.mock as mock

from marvin_product_classifier_engine.prediction import Predictor
from sklearn.naive_bayes import MultinomialNB
import pandas as pd


class TestPredictor:
    def test_execute(self, mocked_params):

        feature_df = pd.DataFrame(data={'col1': [1, 2], 'col2': [3, 4]})
        label_df = pd.DataFrame(data={'col1': [0, 1]})

        clf_tmp = MultinomialNB().fit(feature_df, label_df)
        model_mocked = {"clf": clf_tmp}

        ac = Predictor(model=model_mocked)
        ac.execute(input_message=[1, 2], params=mocked_params)
        assert not ac._params