import os

from qolab.file_utils import infer_compression, save_table_with_header


def test_infer_compression():
    assert infer_compression("data.dat") is None
    assert infer_compression("data.unknown_ext") is None
    assert infer_compression("data.dat.gz") == "gzip"
    assert infer_compression("data.dat.bz") == "bzip"
    assert infer_compression("data.dat.bz2") == "bzip"


def test_save_table_with_header_match_filename_to_compression():
    import tempfile

    data = [[1, 3], [2, 5]]
    headerstr = ["header 1", "header 2"]
    with tempfile.TemporaryDirectory(prefix="qolab_test") as d:
        fname = os.path.join(d, "t.dat")
        assert save_table_with_header(
            fname,
            data,
            header=headerstr,
            compressionmethod=None,
            match_filename_to_compression=True,
        ) == (fname)
        assert save_table_with_header(
            fname,
            data,
            header=headerstr,
            compressionmethod="gzip",
            match_filename_to_compression=True,
        ) == (fname + ".gz")
        assert save_table_with_header(
            fname,
            data,
            header=headerstr,
            compressionmethod="bzip",
            match_filename_to_compression=True,
        ) == (fname + ".bz")

        # now cases when requested extension does not match the compression
        assert save_table_with_header(
            fname + ".gz",
            data,
            header=headerstr,
            compressionmethod=None,
            match_filename_to_compression=True,
        ) == (fname + ".gz.dat")
        assert save_table_with_header(
            fname + ".bz",
            data,
            header=headerstr,
            compressionmethod="gzip",
            match_filename_to_compression=True,
        ) == (fname + ".bz.gz")
        assert save_table_with_header(
            fname + ".gz",
            data,
            header=headerstr,
            compressionmethod="bzip",
            match_filename_to_compression=True,
        ) == (fname + ".gz.bz")

        # do as I told cases, extension does not match compression
        assert save_table_with_header(
            fname + ".gz",
            data,
            header=headerstr,
            compressionmethod=None,
            match_filename_to_compression=False,
        ) == (fname + ".gz")
        assert save_table_with_header(
            fname,
            data,
            header=headerstr,
            compressionmethod="gzip",
            match_filename_to_compression=False,
        ) == (fname)
        assert save_table_with_header(
            fname,
            data,
            header=headerstr,
            compressionmethod="bzip",
            match_filename_to_compression=False,
        ) == (fname)