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')