sklearn/utils/tests/test_user_interface.py - main - githubmirror/scikit-learn

sklearn/utils/tests/test_user_interface.py

65 строк2 KB

Jérémie du Boisberranger

MAINT cleanup utils.__init__: move _print_elapsed_time into dedicated submodule (#28662)

28 мар 2024, 09:52

28 мар 2024, 09:52b599f82

import string
import timeit

import pytest

from sklearn.utils._user_interface import _message_with_time, _print_elapsed_time


@pytest.mark.parametrize(
    ["source", "message", "is_long"],
    [
        ("ABC", string.ascii_lowercase, False),
        ("ABCDEF", string.ascii_lowercase, False),
        ("ABC", string.ascii_lowercase * 3, True),
        ("ABC" * 10, string.ascii_lowercase, True),
        ("ABC", string.ascii_lowercase + "\u1048", False),
    ],
)
@pytest.mark.parametrize(
    ["time", "time_str"],
    [
        (0.2, "   0.2s"),
        (20, "  20.0s"),
        (2000, "33.3min"),
        (20000, "333.3min"),
    ],
)
def test_message_with_time(source, message, is_long, time, time_str):
    out = _message_with_time(source, message, time)
    if is_long:
        assert len(out) > 70
    else:
        assert len(out) == 70

    assert out.startswith("[" + source + "] ")
    out = out[len(source) + 3 :]

    assert out.endswith(time_str)
    out = out[: -len(time_str)]
    assert out.endswith(", total=")
    out = out[: -len(", total=")]
    assert out.endswith(message)
    out = out[: -len(message)]
    assert out.endswith(" ")
    out = out[:-1]

    if is_long:
        assert not out
    else:
        assert list(set(out)) == ["."]


@pytest.mark.parametrize(
    ["message", "expected"],
    [
        ("hello", _message_with_time("ABC", "hello", 0.1) + "\n"),
        ("", _message_with_time("ABC", "", 0.1) + "\n"),
        (None, ""),
    ],
)
def test_print_elapsed_time(message, expected, capsys, monkeypatch):
    monkeypatch.setattr(timeit, "default_timer", lambda: 0)
    with _print_elapsed_time("ABC", message):
        monkeypatch.setattr(timeit, "default_timer", lambda: 0.1)
    assert capsys.readouterr().out == expected