public-engines/kaggle-titanic-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_titanic_engine.prediction import Predictor


class TestPredictor:
    def test_execute(self, mocked_params):

        model_mocked = {
            "rf": mock.MagicMock(),
            "svm": mock.MagicMock()
        }

        ac = Predictor(model=model_mocked)
        ac.execute(input_message="fake message", params=mocked_params)

        model_mocked["rf"].predict.assert_called_once()
        model_mocked["svm"].predict.assert_called_once()