| #!/usr/bin/env python |
| # coding=utf-8 |
| |
| try: |
| import mock |
| |
| except ImportError: |
| import unittest.mock as mock |
| |
| import pandas as pd |
| |
| from marvin_iris_species_engine.data_handler import TrainingPreparator |
| |
| |
| class TestTrainingPreparator: |
| def test_execute(self, mocked_params): |
| mocked_params['test_size'] = 0.3 |
| mocked_params['random_state'] = 10 |
| data = { |
| 'SepalLengthCm': [1, 2], |
| 'SepalWidthCm': [3, 4], |
| 'PetalLengthCm': [5, 6], |
| 'PetalWidthCm': [7, 8], |
| 'Species': 'specie1' |
| } |
| test_dataset = pd.DataFrame(data=data) |
| |
| ac = TrainingPreparator(initial_dataset=test_dataset) |
| ac.execute(params=mocked_params) |