mlflow--mlflow/tests/utils/test_utils.py

import socket
from unittest import mock

import pytest

from mlflow.utils import (
    AttrDict,
    _chunk_dict,
    _get_fully_qualified_class_name,
    _truncate_dict,
    merge_dicts,
)


def test_truncate_dict():
    d = {"12345": "12345"}
    length = 5

    with mock.patch("mlflow.utils._logger.warning") as mock_warning:
        max_length = length - 1

        # Truncate keys
        assert _truncate_dict(d, max_key_length=max_length) == {"1...": "12345"}
        mock_warning.assert_called_once_with("Truncated the key `1...`")
        mock_warning.reset_mock()

        # Truncate values
        assert _truncate_dict(d, max_value_length=max_length) == {"12345": "1..."}
        mock_warning.assert_called_once_with(
            "Truncated the value of the key `12345`. Truncated value: `1...`"
        )
        mock_warning.reset_mock()

        # Truncate both keys and values
        assert _truncate_dict(d, max_key_length=max_length, max_value_length=max_length) == {
            "1...": "1..."
        }
        assert mock_warning.call_count == 2
        (args1, _), (args2, _) = mock_warning.call_args_list
        assert args1[0] == "Truncated the key `1...`"
        assert args2[0] == "Truncated the value of the key `1...`. Truncated value: `1...`"

    assert _truncate_dict(d, max_key_length=length, max_value_length=length) == {"12345": "12345"}
    assert _truncate_dict(d, max_key_length=length + 1, max_value_length=length + 1) == {
        "12345": "12345"
    }

    with pytest.raises(
        ValueError, match="Must specify at least either `max_key_length` or `max_value_length`"
    ):
        _truncate_dict(d)


def test_merge_dicts():
    dict_a = {"a": 3, "b": {"c": {"d": [1, 2, 3]}}, "k": "hello"}
    dict_b = {"test_var": [1, 2]}
    expected_ab = {"a": 3, "b": {"c": {"d": [1, 2, 3]}}, "k": "hello", "test_var": [1, 2]}
    assert merge_dicts(dict_a, dict_b) == expected_ab

    dict_c = {"a": 10}
    with pytest.raises(ValueError, match="contains duplicate keys"):
        merge_dicts(dict_a, dict_c)

    expected_ac = {"a": 10, "b": {"c": {"d": [1, 2, 3]}}, "k": "hello"}
    assert merge_dicts(dict_a, dict_c, raise_on_duplicates=False) == expected_ac


def test_chunk_dict():
    d = {i: i for i in range(10)}
    assert list(_chunk_dict(d, 4)) == [
        {i: i for i in range(4)},
        {i: i for i in range(4, 8)},
        {i: i for i in range(8, 10)},
    ]
    assert list(_chunk_dict(d, 5)) == [
        {i: i for i in range(5)},
        {i: i for i in range(5, 10)},
    ]
    assert list(_chunk_dict(d, len(d))) == [d]
    assert list(_chunk_dict(d, len(d) + 1)) == [d]


def test_get_fully_qualified_class_name():
    class Foo:
        pass

    assert _get_fully_qualified_class_name(Foo()) == f"{__name__}.Foo"


def test_inspect_original_var_name():
    from mlflow.utils import _inspect_original_var_name

    def f1(a1, expected_name):
        assert _inspect_original_var_name(a1, "unknown") == expected_name

    xyz1 = object()
    f1(xyz1, "xyz1")

    f1(str(xyz1), "unknown")

    f1(None, "unknown")

    def f2(b1, expected_name):
        f1(b1, expected_name)

    f2(xyz1, "xyz1")

    def f3(a1, *, b1, expected_a1_name, expected_b1_name):
        assert _inspect_original_var_name(a1, None) == expected_a1_name
        assert _inspect_original_var_name(b1, None) == expected_b1_name

    xyz2 = object()
    xyz3 = object()

    f3(*[xyz2], **{"b1": xyz3, "expected_a1_name": "xyz2", "expected_b1_name": "xyz3"})


def test_random_name_generation():
    from mlflow.utils import name_utils

    # Validate exhausted loop truncation
    name = name_utils._generate_random_name(max_length=8)
    assert len(name) == 8

    # Validate default behavior while calling 1000 times that names end in integer
    names = [name_utils._generate_random_name() for i in range(1000)]
    assert all(len(name) <= 20 for name in names)
    assert all(name[-1].isnumeric() for name in names)


def test_basic_attribute_access():
    d = AttrDict({"a": 1, "b": 2})
    assert d.a == 1
    assert d.b == 2


def test_nested_attribute_access():
    d = AttrDict({"a": 1, "b": {"c": 3, "d": 4}})
    assert d.b.c == 3
    assert d.b.d == 4


def test_non_existent_attribute():
    d = AttrDict({"a": 1, "b": 2})
    with pytest.raises(AttributeError, match="'AttrDict' object has no attribute 'c'"):
        _ = d.c


def test_hasattr():
    d = AttrDict({"a": 1, "b": {"c": 3, "d": 4}})
    assert hasattr(d, "a")
    assert hasattr(d, "b")
    assert not hasattr(d, "e")
    assert hasattr(d.b, "c")
    assert not hasattr(d.b, "e")


def test_subclass_hasattr():
    class SubAttrDict(AttrDict):
        pass

    d = SubAttrDict({"a": 1, "b": {"c": 3, "d": 4}})
    assert hasattr(d, "a")
    assert not hasattr(d, "e")
    assert hasattr(d.b, "c")
    assert not hasattr(d.b, "e")

    with pytest.raises(AttributeError, match="'SubAttrDict' object has no attribute 'g'"):
        _ = d.g


def test_setattr():
    d = AttrDict({"a": 1, "b": 2})

    # Set existing attribute
    d.a = 10
    assert d.a == 10
    assert d["a"] == 10

    # Set new attribute
    d.c = 3
    assert d.c == 3
    assert d["c"] == 3
    assert "c" in d


def test_delattr():
    d = AttrDict({"a": 1, "b": 2, "c": 3})

    # Delete existing attribute
    del d.b
    assert "b" not in d
    assert not hasattr(d, "b")

    # Verify other attributes still exist
    assert d.a == 1
    assert d.c == 3


def test_delattr_non_existent():
    d = AttrDict({"a": 1, "b": 2, "c": 3})
    with pytest.raises(KeyError, match="nonexistent"):
        del d.nonexistent


def test_find_free_port():
    from mlflow.utils import find_free_port

    port = find_free_port()
    assert isinstance(port, int)
    assert 1024 <= port <= 65535


def test_is_port_available_returns_true_for_free_port():
    from mlflow.utils import find_free_port, is_port_available

    port = find_free_port()
    assert is_port_available(port) is True


def test_is_port_available_returns_false_for_bound_port():
    from mlflow.utils import is_port_available

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        bound_port = s.getsockname()[1]
        assert is_port_available(bound_port) is False