from __future__ import annotations

from nxtomomill.models.h52nx._context import (
    extract_contexts,
    Beamline,
    Proxy,
    key_to_str,
    _Key,
)
import pytest


def test_extract_contexts():
    assert extract_contexts("id16a:tomo_config:sample_u") == _Key(
        dataset="sample_u",
        contexts={Proxy.TOMO_CONFIG, Beamline.id16a},
    )
    assert extract_contexts("sample_u") == _Key(
        dataset="sample_u",
        contexts=set(),
    )
    assert extract_contexts("id19:sample_u") == _Key(
        dataset="sample_u",
        contexts={Beamline.id19},
    )

    with pytest.raises(ValueError):
        extract_contexts("id16a:")
    with pytest.raises(ValueError):
        extract_contexts("tomo_config:")
    with pytest.raises(ValueError):
        extract_contexts("unknown:sample_u")


def test_key_to_str():
    assert (
        key_to_str(
            _Key(
                dataset="toto",
                contexts={
                    Beamline.id03,
                    Proxy.TOMO_CONFIG,
                },
            )
        )
        == f"{Beamline.id03.value}:{Proxy.TOMO_CONFIG.value}:toto"
    )
    assert (
        key_to_str(
            _Key(
                dataset="toto",
                contexts={
                    Proxy.TOMO_CONFIG,
                },
            )
        )
        == f"{Proxy.TOMO_CONFIG.value}:toto"
    )
    assert (
        key_to_str(
            _Key(
                dataset="toto",
                contexts=set(),
            )
        )
        == "toto"
    )
    with pytest.raises(TypeError):
        key_to_str(
            {
                "dataset": "toto",
                "contexts": set(),
            }
        )
