test_utils.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import collections
import datetime
import os
import typing as tp
from pathlib import Path

import pydantic
import pytest

import exca

from . import utils
from .confdict import ConfDict
from .utils import to_dict


class BaseModel(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid")


class C(BaseModel):
    param: int = 12
    _exclude_from_cls_uid = (".",)


class A(BaseModel):
    _exclude_from_cls_uid = ("y",)
    x: int = 12
    y: str = "hello"


class B(BaseModel):
    a1: A
    a2: A = A()
    a3: A = A(x=13)
    a4: int = 12
    c: C = C()

    @classmethod
    def _exclude_from_cls_uid(cls) -> tp.List[str]:
        return ["a4"]


def test_to_dict_full() -> None:
    d = to_dict(B(a1={"y": "world"}))  # type: ignore
    out = ConfDict(d).to_yaml()
    expected = """a1:
  x: 12
  y: world
a2:
  x: 12
  y: hello
a3:
  x: 13
  y: hello
a4: 12
c.param: 12
"""
    assert out == expected


def test_to_dict_nondefault() -> None:
    b = B(a1={}, a2={"y": "world"}, a4=13, c={"param": 13})  # type: ignore
    d = to_dict(b, exclude_defaults=True)
    out = ConfDict(d).to_yaml()
    expected = """a1: {}
a2.y: world
a4: 13
c.param: 13
"""
    assert out == expected


def test_to_dict_uid() -> None:
    b = B(a1={}, a2={"y": "world"}, a4=13, c={"param": 13})  # type: ignore
    d = to_dict(b, uid=True, exclude_defaults=True)
    out = ConfDict(d).to_yaml()
    print(out)
    expected = "a1: {}\n"
    assert out == expected


class D2(BaseModel):
    uid: tp.Literal["D2"] = "D2"


class D1(BaseModel):
    uid: tp.Literal["D1"] = "D1"
    anything: int = 12
    sub: D2 = D2()


class Discrim(BaseModel):
    inst: D1 | D2 = pydantic.Field(..., discriminator="uid")
    something_else: tp.List[str] | int
    seq: tp.List[tp.List[tp.Annotated[D1 | D2, pydantic.Field(discriminator="uid")]]]
    stuff: tp.List[D1] = []


def test_missing_discriminator() -> None:
    class DiscrimD(BaseModel):
        instd: D1 | D2

    _ = DiscrimD(instd={"uid": "D1"})  # type: ignore


def test_discriminators(caplog: tp.Any) -> None:
    d = Discrim(
        inst={"uid": "D2"},  # type: ignore
        something_else=12,
        seq=[[{"uid": "D2"}, {"uid": "D1"}]],  # type: ignore
    )
    expected = """inst.uid: D2
seq:
- - uid: D2
  - anything: 12
    sub.uid: D2
    uid: D1
something_else: 12
stuff: []
"""
    # check uid of subinstance (should not have discriminator)
    sub_out = ConfDict.from_model(d.inst, exclude_defaults=True)
    assert not sub_out
    # check uid of instance (should have discriminators)
    out = ConfDict.from_model(d).to_yaml()
    assert out == expected
    expected = """inst.uid: D2
seq:
- - uid: D2
  - uid: D1
something_else: 12
"""
    out = ConfDict.from_model(d, exclude_defaults=True).to_yaml()
    assert not caplog.records
    assert out == expected
    # check uid of subinstance again (should not have discriminators)
    sub_out = ConfDict.from_model(d.inst, exclude_defaults=True)
    assert not sub_out
    # CHECK AGAIN THE FULL STUFF!
    out = ConfDict.from_model(d, exclude_defaults=True).to_yaml()
    assert out == expected


def test_recursive_freeze() -> None:
    d = Discrim(
        inst={"uid": "D2"},  # type: ignore
        something_else=12,
        seq=[[{"uid": "D2"}, {"uid": "D1"}]],  # type: ignore
    )
    sub = d.seq[0][0]
    with pytest.raises(ValueError):
        # not frozen but field does not exist
        sub.blublu = 12  # type: ignore
    utils.recursive_freeze(d)
    if hasattr(sub, "_setattr_handler"):
        with pytest.raises(RuntimeError):
            # frozen, otherwise it would be a value error
            sub.blublu = 12  # type: ignore
    else:
        assert sub.model_config["frozen"]


class OptDiscrim(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid")
    val: tp.Annotated[D1 | D2, pydantic.Field(discriminator="uid")] | None = None


def test_optional_discriminator(caplog: tp.Any) -> None:
    d = OptDiscrim(val={"uid": "D2"})  # type: ignore
    out = ConfDict.from_model(d, exclude_defaults=True).to_yaml()
    assert not caplog.records
    expected = "val.uid: D2\n"
    assert out == expected


class RecursiveLeaf(BaseModel):
    edge_type: tp.Literal["leaf"] = "leaf"


class RecursiveEdge(BaseModel):
    infra: exca.TaskInfra = exca.TaskInfra(cluster=None)
    edge_type: tp.Literal["edge"] = "edge"
    child: "RecursiveElem"

    @infra.apply
    def run(self) -> None:
        return


RecursiveElem = tp.Annotated[
    RecursiveLeaf | RecursiveEdge,
    pydantic.Field(discriminator="edge_type"),
]


def test_recursive_discriminated_union(tmp_path: Path) -> None:
    """Test that recursive discriminated unions work with infra.config().

    This tests the fix for KeyError when using recursive types with
    discriminated unions, where model_json_schema() returns a $ref schema.
    """
    cfg = RecursiveEdge(
        infra=exca.TaskInfra(cluster=None, folder=tmp_path), child=RecursiveLeaf()
    )
    # This should work without KeyError
    result = cfg.infra.config(exclude_defaults=True)
    assert result == {"child": {"edge_type": "leaf"}}
    # Test with uid=True as well
    result_uid = cfg.infra.config(uid=True, exclude_defaults=True)
    assert "child" in result_uid
    # Test deep recursion
    cfg_deep = RecursiveEdge(
        infra=exca.TaskInfra(cluster=None, folder=tmp_path),
        child=RecursiveEdge(
            infra=exca.TaskInfra(cluster=None, folder=tmp_path), child=RecursiveLeaf()
        ),
    )
    result_deep = cfg_deep.infra.config(exclude_defaults=True)
    assert result_deep == {"child": {"child": {"edge_type": "leaf"}, "edge_type": "edge"}}


@pytest.mark.parametrize("replace", (True, False))
@pytest.mark.parametrize("existing_content", [None, "blublu"])
def test_temporary_save_path(
    tmp_path: Path, existing_content: str | None, replace: bool
) -> None:
    filepath = tmp_path / "save_and_move_test.txt"
    if existing_content:
        filepath.write_text(existing_content)
    with utils.temporary_save_path(filepath, replace=replace) as tmp:
        assert str(tmp).endswith(".txt")
        tmp.write_text("12")
        if existing_content:
            assert filepath.read_text("utf8") == existing_content
    expected = "12"
    if existing_content is not None and not replace:
        expected = "blublu"
    assert filepath.read_text("utf8") == expected


def test_temporary_save_path_error() -> None:
    with pytest.raises(FileNotFoundError):
        with utils.temporary_save_path("save_and_move_test"):
            pass


@pytest.mark.parametrize(
    "hint,expected",
    [
        (None | int, []),
        (None | D1, [D1]),
        (D2 | D1, [D2, D1]),
        (D1, [D1]),
        (list[D2 | D1], [D2, D1]),
        (
            tp.List[tp.List[tp.Annotated[D1 | D2, pydantic.Field(discriminator="uid")]]],
            [D1, D2],
        ),
        (tp.Annotated[D1 | D2, pydantic.Field(discriminator="uid")] | None, [D1, D2]),  # type: ignore
    ],
)
def test_pydantic_hints(hint: tp.Any, expected: tp.List[tp.Any]) -> None:
    assert tuple(utils._pydantic_hints(hint)) == tuple(expected)


def test_environment_variable_context() -> None:
    name = "ENV_VAR_TEST"
    assert name not in os.environ
    with utils.environment_variables(ENV_VAR_TEST="blublu"):
        assert os.environ[name] == "blublu"
        with utils.environment_variables(ENV_VAR_TEST="blublu2"):
            assert os.environ[name] == "blublu2"
        assert os.environ[name] == "blublu"
    assert name not in os.environ


def test_iter_string_values():
    out = dict(utils._iter_string_values({"a": [12, {"b": 13, "c": "val"}]}))
    assert out == {"a.1.c": "val"}


class MissingForbid(pydantic.BaseModel):
    param: int = 12


class WithMissingForbid(BaseModel):
    missing: MissingForbid = MissingForbid()


def test_extra_forbid() -> None:
    m = MissingForbid()
    with pytest.raises(RuntimeError):
        ConfDict.from_model(m, uid=True, exclude_defaults=True)
    w = WithMissingForbid()
    with pytest.raises(RuntimeError):
        ConfDict.from_model(w, uid=True, exclude_defaults=True)


class D(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid")
    x: int = 12


class A12(BaseModel):
    _exclude_from_cls_uid = ("y",)
    name: str = "name"
    unneeded: str = "is default"
    x: int = 12
    y: str = "hello"


class NewDefault(BaseModel):
    a: A12 = A12(x=13)


@pytest.mark.parametrize("with_y", (False, True))
@pytest.mark.parametrize(
    "value,expected",
    [
        (11, "a.x: 11"),
        (12, "a: {}"),
        (13, "{}"),
    ],
)
def test_new_default(value: int, expected: str, with_y: bool) -> None:
    params: tp.Any = {"x": value}
    if with_y:
        params["y"] = "world"
    m = NewDefault(a=params)
    out = ConfDict.from_model(m, uid=True, exclude_defaults=True)
    assert out.to_yaml().strip() == expected
    m2 = NewDefault(**out)
    assert m2.a.x == value


class NewDefaultOther(BaseModel):
    a: A12 = A12(x=13, y="stuff")


def test_new_default_other() -> None:
    m = NewDefaultOther(a={"x": 13})  # type: ignore
    out = ConfDict.from_model(m, uid=True, exclude_defaults=True)
    assert out.to_yaml().strip() == "{}"


class NewDefaultOther2diff(BaseModel):
    a: A12 = A12(x=13, unneeded="something else", y="stuff")


def test_new_default_other2diff() -> None:
    # revert unneeded to default, so it wont show in model_dump, but we need to define x=13
    m = NewDefaultOther2diff(a={"x": 13, "unneeded": "is default"})  # type: ignore
    out = ConfDict.from_model(m, uid=True, exclude_defaults=True)
    assert out.to_yaml().strip() == "a.x: 13"


class ActualDefaultOverride(BaseModel):
    a: A12 = A12(x=12)
    a_default: A12 = A12()


def test_actual_default_override() -> None:
    m = ActualDefaultOverride(a={"x": 13})  # type: ignore
    out = ConfDict.from_model(m, uid=True, exclude_defaults=True)
    assert out.to_yaml().strip() == "a.x: 13"
    #
    m = ActualDefaultOverride(a={"x": 12, "y": "stuff"}, a_default={"x": 12, "y": "stuff"})  # type: ignore
    out = ConfDict.from_model(m, uid=True, exclude_defaults=True)
    assert out.to_yaml().strip() == "{}"


class DiscrimDump(BaseModel):
    inst: D1 | D2 = pydantic.Field(D1(), discriminator="uid")


def test_dump() -> None:
    dd = DiscrimDump(inst={"uid": "D1"})  # type: ignore
    out = ConfDict.from_model(dd, uid=True, exclude_defaults=True)
    assert not out
    dd = DiscrimDump(inst={"uid": "D2"})  # type: ignore
    out = ConfDict.from_model(dd, uid=True, exclude_defaults=True)
    assert out == {"inst": {"uid": "D2"}}


D1D2 = tp.Annotated[D1 | D2, pydantic.Field(discriminator="uid")]


class OrderedDump(BaseModel):
    insts: collections.OrderedDict[str, D1D2] = collections.OrderedDict()


def test_ordered_dict() -> None:
    od = OrderedDump(insts={"blublu": {"uid": "D1"}, "stuff": {"uid": "D2"}, "blublu2": {"uid": "D1"}})  # type: ignore
    out = ConfDict.from_model(od, uid=True, exclude_defaults=True)
    # check that nothing alters the order
    assert isinstance(out["insts"], collections.OrderedDict)
    assert tuple(out["insts"].keys()) == ("blublu", "stuff", "blublu2")
    out["insts.blublu.anything"] = 144
    assert tuple(out["insts"].keys()) == ("blublu", "stuff", "blublu2")
    out["insts.blublu2.anything"] = 144
    assert tuple(out["insts"].keys()) == ("blublu", "stuff", "blublu2")
    assert isinstance(out["insts"], collections.OrderedDict)
    # keys should be ordered in name and hash:
    uid = "insts={blublu={uid=D1,anything=144},stuff.uid=D2,blublu2={uid=D1,anything=144}}-46863fcc"
    assert out.to_uid() == uid


class ComplexDiscrim(BaseModel):
    inst: dict[str, tuple[D1D2, bool]] | None = None


def test_complex_discrim() -> None:
    d = ComplexDiscrim(inst={"stuff": ({"uid": "D2"}, True)})  # type: ignore
    out = ConfDict.from_model(d, uid=True, exclude_defaults=True)
    assert utils.DISCRIMINATOR_FIELD in d.inst["stuff"][0].__dict__  # type: ignore
    assert "D2" in out.to_uid()


class HierarchicalCfg(pydantic.BaseModel):
    a: A = A()
    _a: A = A()
    c: C = C()
    content: tp.List["HierarchicalCfg"] = []


def test_find_models() -> None:
    hcfg = HierarchicalCfg(content=[{}, {}])  # type: ignore
    out = utils.find_models(hcfg, A)
    assert set(out) == {
        "a",
        "content.0.a",
        "content.1.a",
        "_a",
        "content.0._a",
        "content.1._a",
    }
    assert all(isinstance(y, A) for y in out.values())


def test_fast_unlink(tmp_path: Path) -> None:
    # file
    fp = tmp_path / "blublu.txt"
    fp.touch()
    assert fp.exists()
    with utils.fast_unlink(fp):
        pass
    assert not fp.exists()
    # folder
    fp = tmp_path / "blublu"
    fp.mkdir()
    (fp / "stuff.txt").touch()
    with utils.fast_unlink(fp):
        pass
    assert not fp.exists()


class ComplexTypesConfig(BaseModel):
    x: pydantic.DirectoryPath = Path("/")
    y: datetime.timedelta = datetime.timedelta(minutes=1)
    # ImportString needs to work (serialize_as_any=False in model_dump)
    z: pydantic.ImportString = ConfDict


def test_complex_types() -> None:
    c = ComplexTypesConfig()
    out = ConfDict.from_model(c, uid=True, exclude_defaults=False)
    expected = """x: /
y: PT1M
z: exca.confdict.ConfDict
"""
    assert out.to_yaml() == expected
    assert out.to_uid().startswith("x=-,y=PT1M,z=exca.confdict.ConfDict")


class BasicP(pydantic.BaseModel):
    b: pydantic.BaseModel | None = None
    infra: exca.TaskInfra = exca.TaskInfra(version="12")

    @infra.apply
    def func(self) -> int:
        return 12


def test_basic_pydantic() -> None:
    b = BasicP(b={"uid": "D2"})  # type: ignore
    with pytest.raises(RuntimeError) as e:
        b.infra.clone_obj()
    assert "discriminated union" in e.value.args[0]


class CO(BaseModel):
    stuff: str = "blublu"

    def _exca_uid_dict_override(self) -> dict[str, tp.Any]:
        return {"override": "success"}


class ConfWithOverride(BaseModel):
    a: A = A()
    s: CO = CO()


@pytest.mark.parametrize("uid", (True, False))
@pytest.mark.parametrize("exc", (True, False))
@pytest.mark.parametrize("raw", (True, False))
@pytest.mark.parametrize("bypass", (True, False))
@pytest.mark.parametrize("use_exporter", (True, False))
def test_uid_dict_override(
    uid: bool, exc: bool, raw: bool, bypass: bool, use_exporter: bool
) -> None:
    # use model with override as model or sub-model
    if raw:
        model = CO(stuff="blu")
    else:
        model = ConfWithOverride(s={"stuff": "blu"})  # type: ignore
    # use the ConfDict directly, or the exporter (which allows bypassing the override)
    if use_exporter:
        exporter = utils.ConfigExporter(
            uid=uid, exclude_defaults=exc, ignore_first_override=bypass
        )
        cfg = ConfDict(exporter.apply(model))
    else:
        cfg = ConfDict.from_model(model, uid=uid, exclude_defaults=exc)
    out = cfg.to_yaml()
    if uid and exc and not (use_exporter and raw and bypass):
        assert "override" in out
    else:
        assert "override" not in out


def test_check_configs(tmp_path: Path) -> None:
    """Test ConfigDump: creates files, detects inconsistencies, handles corruption."""
    model = A(x=5, y="test")
    dump = utils.ConfigDump(model=model)

    # Creates config files
    dump.check_and_write(tmp_path)
    assert (tmp_path / "uid.yaml").exists()
    assert "x: 5" in (tmp_path / "uid.yaml").read_text("utf8")

    # Detects inconsistent uid.yaml (error includes model info)
    (tmp_path / "uid.yaml").write_text("x: 999\n")
    with pytest.raises(RuntimeError, match="(?s)Inconsistent.*this is for object"):
        dump.check_and_write(tmp_path)

    # Handles corrupted files (deletes and recreates)
    (tmp_path / "uid.yaml").write_text("invalid: yaml: {{")
    dump.check_and_write(tmp_path)
    assert "x: 5" in (tmp_path / "uid.yaml").read_text("utf8")

    # Works with list of models (writes as list, not wrapped)
    models = [A(x=1), A(x=2)]
    folder2 = tmp_path / "list_test"
    folder2.mkdir()
    utils.ConfigDump(model=models).check_and_write(folder2)
    content = (folder2 / "uid.yaml").read_text("utf8")
    assert content == "- x: 1\n- x: 2\n"


def test_permission_setter(tmp_path: Path) -> None:
    # Test with permissions enabled
    ps = utils.PermissionSetter(0o777)
    folder = tmp_path / "test_folder"
    folder.mkdir()
    ps.set(folder)
    assert oct(folder.stat().st_mode)[-3:] == "777"

    # Test with file
    fp = folder / "test.txt"
    fp.write_text("hello")
    ps.set(fp)
    assert oct(fp.stat().st_mode)[-3:] == "777"

    # Test with permissions disabled (None)
    ps_none = utils.PermissionSetter(None)
    folder2 = tmp_path / "test_folder2"
    folder2.mkdir()
    before_mode = folder2.stat().st_mode
    ps_none.set(folder2)  # should do nothing
    assert folder2.stat().st_mode == before_mode


def test_permission_setter_recursive(tmp_path: Path) -> None:
    ps = utils.PermissionSetter(0o777)

    # Create nested structure
    root = tmp_path / "root"
    sub = root / "sub"
    sub.mkdir(parents=True)
    fp = sub / "file.txt"
    fp.write_text("content")

    # Set permissions recursively
    ps.set(root, recursive=True)
    assert oct(root.stat().st_mode)[-3:] == "777"
    assert oct(sub.stat().st_mode)[-3:] == "777"
    assert oct(fp.stat().st_mode)[-3:] == "777"