test_format.py

# -*- coding: utf-8 -*-

from pathlib import Path

import pytest

from dan.datasets.format.atr import ATRDatasetFormatter, remove_spaces


@pytest.mark.parametrize(
    "text,trimmed",
    (
        ("no_spaces", "no_spaces"),
        (" beginning", "beginning"),
        ("ending ", "ending"),
        (" both ", "both"),
        ("    consecutive", "consecutive"),
        ("\ttab", "tab"),
    ),
)
def test_remove_spaces(text, trimmed):
    assert remove_spaces(text) == trimmed


@pytest.mark.parametrize(
    "image_format, expected_format",
    (
        (".png", "png"),
        ("png", "png"),
    ),
)
def test_image_format_parsing(image_format, expected_format):
    formatter = ATRDatasetFormatter(
        dataset=Path("my_dataset"), image_format=image_format, remove_spaces=True
    )
    assert formatter.image_format == expected_format