import os
import os.path

import pytest

from psyrun.core import load_infile, load_results, save_outfile
from psyrun.pspace import Param
from psyrun.split import Splitter


@pytest.mark.parametrize(
    'pspace_size,max_splits,min_items,n_splits', [
        (7, 4, 4, 2), (8, 4, 4, 2), (9, 4, 4, 3),
        (15, 4, 4, 4), (16, 4, 4, 4), (17, 4, 4, 4),
        (15, 2, 4, 2), (16, 4, 16, 1)
    ])
class TestSplitter(object):
    def test_n_splits(
            self, tmpdir, pspace_size, max_splits, min_items, n_splits):
        splitter = Splitter(
            str(tmpdir), Param(x=range(pspace_size)), max_splits, min_items)
        assert splitter.n_splits == n_splits
        assert len(list(splitter.iter_in_out_files())) == n_splits

    def test_split_merge_roundtrip(
            self, tmpdir, pspace_size, max_splits, min_items, n_splits):
        splitter = Splitter(
            str(tmpdir), Param(x=range(pspace_size)), max_splits, min_items)
        splitter.split()

        for filename in os.listdir(splitter.indir):
            infile = os.path.join(splitter.indir, filename)
            outfile = os.path.join(splitter.outdir, filename)
            save_outfile(load_infile(infile), outfile)

        result_file = os.path.join(str(tmpdir), 'result.h5')
        Splitter.merge(str(tmpdir), result_file)
        result = load_results(result_file)
        assert sorted(result['x']) == sorted(range(pspace_size))

Add unit test for Splitter.
Add unit test for Splitter.
