60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
from model_manager.sientia import utils
|
|
|
|
|
|
def test_split_train_test_default(monkeypatch):
|
|
captured_args = {}
|
|
|
|
def fake_train_test_split(*args, **kwargs):
|
|
captured_args['args'] = args
|
|
captured_args['kwargs'] = kwargs
|
|
return ('X_train', 'X_test', 'y_train', 'y_test')
|
|
|
|
monkeypatch.setattr(utils, 'train_test_split', fake_train_test_split)
|
|
|
|
X = [1, 2, 3, 4]
|
|
y = [0, 1, 0, 1]
|
|
|
|
result = utils.split_train_test(X, y)
|
|
|
|
assert captured_args['args'] == (X, y)
|
|
assert captured_args['kwargs'] == {
|
|
'test_size': None,
|
|
'train_size': None,
|
|
'random_state': None,
|
|
'shuffle': True,
|
|
'stratify': None,
|
|
}
|
|
assert result == ('X_train', 'X_test', 'y_train', 'y_test')
|
|
|
|
|
|
def test_split_train_test_with_parameters(monkeypatch):
|
|
captured_kwargs = {}
|
|
|
|
def fake_train_test_split(*args, **kwargs):
|
|
captured_kwargs.update(kwargs)
|
|
return ('train_X', 'test_X', 'train_y', 'test_y')
|
|
|
|
monkeypatch.setattr(utils, 'train_test_split', fake_train_test_split)
|
|
|
|
X = [[1], [2], [3], [4]]
|
|
y = [0, 1, 0, 1]
|
|
|
|
result = utils.split_train_test(
|
|
X,
|
|
y,
|
|
test_size=0.25,
|
|
train_size=0.75,
|
|
random_state=42,
|
|
shuffle=False,
|
|
stratify=y,
|
|
)
|
|
|
|
assert captured_kwargs == {
|
|
'test_size': 0.25,
|
|
'train_size': 0.75,
|
|
'random_state': 42,
|
|
'shuffle': False,
|
|
'stratify': y,
|
|
}
|
|
assert result == ('train_X', 'test_X', 'train_y', 'test_y')
|