ha-nvidia-shield/tests/test_config_flow_protocol.py

"""Phase 2 step 2 tests: protocol cert/PIN pairing, reauth, entry migration.

Self-contained scoped-stub variant of the config-flow test harness (same
pattern as test_config_flow_pairing.py, with pair-hygiene capabilities added:
FakeHass records unique_id sets, entry updates and reload scheduling). The
androidtvremote2 client is faked with a deterministic pairing code and MAC.
"""

import asyncio  # noqa: F401  (async steps need the event loop)
import sys
import types
from types import SimpleNamespace

import entity_stublib
import pytest
from conftest import _SHIELD_DIR, _load_module
from nsh.const import (
    CONF_ADB_ENABLED,
    CONF_ADB_HOST,
    CONF_ADB_PORT,
    CONF_ADB_STORE_APP_INFO,
    CONF_HOST,
    CONF_JELLYFIN_ENTITY,
    CONF_PAIRING_CODE,
    CONF_POWER_ENTITY,
    CONF_PROTOCOL_ENABLED,
    CONF_PROTOCOL_HOST,
    CONF_RECEIVER_ENTITY,
    CONF_RECEIVER_INPUT,
    DEFAULT_ADB_PORT,
)

_STUB_NAMES = (
    "voluptuous",
    "homeassistant",
    "homeassistant.config_entries",
    "homeassistant.core",
    "homeassistant.const",
    "homeassistant.helpers",
    "homeassistant.helpers.entity_registry",
    "homeassistant.helpers.device_registry",
    "homeassistant.helpers.selector",
    "nsh.helpers",
    "nsh.config_flow",
    "nsh.entry",
)

FAKE_PAIRING_CODE = "8421"
FAKE_DEVICE_NAME = "mdarcy"
FAKE_MAC = "aa:bb:cc:dd:ee:ff"


class FakeInvalidAuthError(Exception):
    """androidtvremote2.exceptions.InvalidAuth stand-in."""


class FakeProtocolClient:
    """Proto client used by the flow pairing path (deterministic)."""

    def __init__(self, fail_code=False):
        self.cert_generated = False
        self.pairing_started = False
        self.finished_codes: list[str] = []
        self.disconnected = 0
        self._fail_code = fail_code

    async def async_generate_cert_if_missing(self):
        self.cert_generated = True
        return True

    async def async_start_pairing(self):
        self.pairing_started = True
        return FAKE_PAIRING_CODE  # deterministic "device-side PIN"

    async def async_finish_pairing(self, pairing_code: str) -> None:
        if self._fail_code or pairing_code != FAKE_PAIRING_CODE:
            raise FakeInvalidAuthError("device rejected the pairing code")
        self.finished_codes.append(pairing_code)

    async def async_get_name_and_mac(self):
        return FAKE_DEVICE_NAME, FAKE_MAC

    def disconnect(self):
        self.disconnected += 1


def _install_stubs() -> None:
    """Import-time stand-ins (idempotent, scoped by the autouse fixture)."""
    if sys.modules.get("voluptuous") is not None:
        return

    vol = types.ModuleType("voluptuous")
    vol.Required = lambda key, default=None: key
    vol.Optional = lambda key, default=None: key
    vol.Schema = lambda schema: schema
    sys.modules["voluptuous"] = vol

    ha = types.ModuleType("homeassistant")
    sys.modules["homeassistant"] = ha

    config_entries = types.ModuleType("homeassistant.config_entries")

    class ConfigFlow:
        VERSION = 1

        def __init_subclass__(cls, **kwargs):
            kwargs.pop("domain", None)
            super().__init_subclass__(**kwargs)

        def async_show_form(self, *, step_id, data_schema=None, errors=None):
            return _form(step_id, errors or {})

        def async_show_menu(self, *, step_id, menu_options, sort=False,
                            description_placeholders=None, errors=None):
            # t56: the real MENU result shape; the flow manager dispatches the
            # chosen `next_step_id` to `async_step_<choice>(None)`.
            return {
                "type": "menu",
                "step_id": step_id,
                "menu_options": list(menu_options),
                "errors": dict(errors or {}),
            }

        def async_create_entry(self, *, title, data, options=None):
            return {"type": "create_entry", "title": title, "data": data,
                    "options": options}

        def async_abort(self, *, reason):
            return {"type": "abort", "reason": reason}

        async def async_set_unique_id(self, unique_id):
            self._unique_id = unique_id

        def _abort_if_unique_id_configured(self, *, error=None, **kwargs):
            if getattr(self, "_unique_id", None) in active_unique_ids:
                raise AbortIfConfiguredError()

    def _form(step_id, errors):
        return {"type": "form", "step_id": step_id, "errors": dict(errors)}

    class AbortIfConfiguredError(Exception):
        pass

    active_unique_ids = set()

    class OptionsFlow:
        pass
    # t29 G3b: the current options-flow API surface on the local stub
    class _SuggestedMixin:
        def add_suggested_values_to_schema(self, schema, suggested):
            schema_dict = getattr(schema, "schema", None) or (
                schema if isinstance(schema, dict) else {}
            )
            schema_dict.setdefault("__suggested__", dict(suggested))
            return schema

    class OptionsFlowWithReload(_SuggestedMixin):
        def async_create_entry(self, *, title, data):
            entry_id = getattr(self.config_entry, "entry_id", None)
            scheduler = getattr(getattr(self, "hass", None), "config_entries", None)
            if entry_id is not None and scheduler is not None and hasattr(
                scheduler, "async_schedule_reload"
            ):
                scheduler.async_schedule_reload(entry_id)
            return {"type": "create_entry", "title": title, "data": data}



    # t29 G3a: the current config-entries API surface on the local stub
    ConfigFlow._get_reauth_entry = lambda self, entry=None: entry

    def _update_reload_and_abort(self, entry, **kwargs):
        # `options` REPLACES the entry options (the real signature); the merge
        # form is still accepted for older callers.
        updates = dict(kwargs.get("options_updates") or {})
        updates.update(kwargs.get("options") or {})
        options = entry.options
        for key, value in updates.items():
            options[key] = value
        if kwargs.get("unique_id"):
            entry.unique_id = kwargs["unique_id"]
        self.hass.config_entries.async_schedule_reload(entry.entry_id)
        reason = kwargs.get("reason")
        if reason is None:
            reason = (
                "reconfigure_successful"
                if getattr(self, "source", None) == "reconfigure"
                else "reauth_successful"
            )
        return self.async_abort(reason=reason)

    ConfigFlow.async_update_reload_and_abort = _update_reload_and_abort

    config_entries.ConfigFlow = ConfigFlow
    config_entries.OptionsFlow = OptionsFlow
    config_entries.OptionsFlowWithReload = OptionsFlowWithReload
    config_entries.ConfigEntry = type("ConfigEntry", (), {})
    sys.modules["homeassistant.config_entries"] = config_entries

    core = types.ModuleType("homeassistant.core")
    core.HomeAssistant = object
    core.callback = lambda fn: fn
    sys.modules["homeassistant.core"] = core

    ha_const = types.ModuleType("homeassistant.const")
    ha_const.EVENT_HOMEASSISTANT_STOP = "homeassistant_stop"
    class _Platform:
        MEDIA_PLAYER = "media_player"
        SENSOR = "sensor"
        REMOTE = "remote"
        BUTTON = "button"
    ha_const.Platform = _Platform
    sys.modules["homeassistant.const"] = ha_const

    helpers = types.ModuleType("homeassistant.helpers")
    entity_registry = types.ModuleType("homeassistant.helpers.entity_registry")

    class _Registry:
        entities: dict = {}

        def async_get(self, entity_id):
            return None

    entity_registry.async_get = lambda hass: _Registry()
    helpers.entity_registry = entity_registry
    sys.modules["homeassistant.helpers.entity_registry"] = entity_registry
    sys.modules["homeassistant.helpers"] = helpers

    device_registry = types.ModuleType("homeassistant.helpers.device_registry")
    device_registry.async_get = lambda hass: SimpleNamespace(devices={})
    helpers.device_registry = device_registry
    sys.modules["homeassistant.helpers.device_registry"] = device_registry
    selector = types.ModuleType("homeassistant.helpers.selector")
    entity_stublib.install_selector_faces(selector)
    sys.modules["homeassistant.helpers.selector"] = selector


class FakeConfigEntries:
    def __init__(self):
        self.updating_entry_calls: list = []
        self.reloaded_entry_ids: list[str] = []

    def async_get_entry(self, entry_id):
        return getattr(self, f"__entry_{entry_id}", None)

    def async_update_entry(self, entry, **kwargs):
        self.updating_entry_calls.append((entry, kwargs))
        for key, value in kwargs.items():
            if key in ("options", "data"):
                setattr(entry, key, value)
            else:
                setattr(entry, key, value)

    def async_schedule_reload(self, entry_id):
        self.reloaded_entry_ids.append(entry_id)


class FakeHass:
    def __init__(self, tmp_dir: str):
        import types as t

        self.config = t.SimpleNamespace(config_dir=tmp_dir)
        self.config_entries = FakeConfigEntries()

    @staticmethod
    async def async_add_executor_job(fn, *args):
        return fn(*args)


def _form_module():
    _install_stubs()
    if "nsh.helpers" not in sys.modules:
        _load_module("nsh.helpers", _SHIELD_DIR / "helpers.py")
    if "nsh.config_flow" not in sys.modules:
        _load_module("nsh.config_flow", _SHIELD_DIR / "config_flow.py")
    return sys.modules["nsh.config_flow"]


def _entry_module():
    _install_stubs()
    if "nsh.entry" not in sys.modules:
        _load_module("nsh.entry", _SHIELD_DIR / "__init__.py")
    return sys.modules["nsh.entry"]


@pytest.fixture(autouse=True)
def scoped_stubs():
    """Keep the import stubs inside this module's tests (t7 discipline)."""
    snapshot = entity_stublib.snapshot_modules()
    _install_stubs()
    yield
    entity_stublib.restore_modules(snapshot)


def make_flow(tmp_dir: str, monkeypatch: pytest.MonkeyPatch, client: FakeProtocolClient):
    """ConfigFlow with the real client factory mocked to our fake."""
    module = _form_module()
    flow = module.NvidiaShieldConfigFlow()
    flow.hass = FakeHass(tmp_dir)
    flow._user_data = {
        "name": "mdarcy",
        CONF_JELLYFIN_ENTITY: "",
        CONF_RECEIVER_ENTITY: "",
        CONF_RECEIVER_INPUT: "",
        CONF_POWER_ENTITY: "",
        CONF_ADB_ENABLED: False,
        CONF_ADB_STORE_APP_INFO: True,
        CONF_HOST: "shield.local",  # t41: ONE address field
        CONF_ADB_PORT: DEFAULT_ADB_PORT,
        CONF_PROTOCOL_ENABLED: True,
    }
    from nsh.connection import ShieldConnectionManager

    monkeypatch.setattr(
        ShieldConnectionManager,
        "build_protocol_client",
        lambda *args, **kwargs: client,
    )
    monkeypatch.setattr(
        module,
        "protocol_cert_paths",
        lambda hass: ("/tmp/nsh-proto-test.pem", "/tmp/nsh-proto-test.key"),
    )
    return flow


class TestProtocolPairing:
    async def test_first_visit_starts_pairing_and_shows_code_form(self, tmp_path, monkeypatch):
        client = FakeProtocolClient()
        flow = make_flow(str(tmp_path), monkeypatch, client)

        result = await flow.async_step_protocol_pair(None)

        assert client.cert_generated is True
        assert client.pairing_started is True
        assert result["type"] == "form"
        assert result["step_id"] == "protocol_pair"
        assert result["errors"] == {}

    async def test_submit_finishes_with_code_and_records_mac_unique_id(self, tmp_path, monkeypatch):
        client = FakeProtocolClient()
        flow = make_flow(str(tmp_path), monkeypatch, client)

        await flow.async_step_protocol_pair(None)
        result = await flow.async_step_protocol_pair(
            {CONF_PAIRING_CODE: FAKE_PAIRING_CODE}
        )

        assert client.finished_codes == [FAKE_PAIRING_CODE]
        assert flow._unique_id == FAKE_MAC
        assert client.disconnected == 1
        # t56: a finished pairing hands over to the setup menu; its Finish row
        # creates the entry
        assert result["type"] == "menu"
        result = await flow.async_step_finish(None)
        assert result["type"] == "create_entry"
        assert result["options"][CONF_PROTOCOL_ENABLED] is True
        assert result["options"][CONF_HOST] == "shield.local"

    async def test_wrong_code_asks_again_then_succeeds(self, tmp_path, monkeypatch):
        client = FakeProtocolClient()
        flow = make_flow(str(tmp_path), monkeypatch, client)

        await flow.async_step_protocol_pair(None)
        failed = await flow.async_step_protocol_pair({"pairing_code": "9999"})
        assert failed["type"] == "form"
        assert failed["errors"] == {"base": "protocol_pair_failed"}

        result = await flow.async_step_protocol_pair(
            {CONF_PAIRING_CODE: FAKE_PAIRING_CODE}
        )
        assert result["type"] == "menu"
        result = await flow.async_step_finish(None)
        assert result["type"] == "create_entry"
        assert client.finished_codes == [FAKE_PAIRING_CODE]


class TestReauthFlow:
    async def test_reauth_repairs_entry_and_schedules_reload(self, tmp_path, monkeypatch):
        client = FakeProtocolClient()
        flow = make_flow(str(tmp_path), monkeypatch, client)
        entry = SimpleNamespace(
            entry_id="e1", title="mdarcy",
            options={CONF_PROTOCOL_ENABLED: True, CONF_HOST: "shield.local"},
            data={},
        )
        flow.hass.config_entries.__dict__["__entry_e1"] = entry
        flow.context = {"entry_id": "e1"}

        shown = await flow.async_step_reauth({CONF_ADB_PORT: DEFAULT_ADB_PORT})
        assert shown["type"] == "form"
        assert shown["step_id"] == "reauth_confirm"

        pair_form = await flow.async_step_reauth_confirm({"confirm": True})
        assert pair_form["step_id"] == "protocol_pair"  # PIN form open on device

        result = await flow.async_step_protocol_pair(
            {CONF_PAIRING_CODE: FAKE_PAIRING_CODE}
        )
        assert result == {"type": "abort", "reason": "reauth_successful"}
        assert entry.options[CONF_PROTOCOL_ENABLED] is True
        assert flow.hass.config_entries.reloaded_entry_ids == ["e1"]


class TestEntryMigration:
    """t41: v2 → v3 collapses the two address fields into CONF_HOST.

    The protocol host wins, else the ADB host (both name the same Shield); the
    legacy keys are removed so nothing can read a stale duplicate. A v1 entry
    travels the whole way in one call.
    """

    async def test_v1_keeps_its_adb_options_and_gains_both_defaults(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(
            version=1,
            data={},
            options={
                "adb_enabled": True,
                "adb_port": 5555,
                "jellyfin_entity": "jf.x",
                "adb_host": "10.0.0.31",  # legacy address
            },
            title="mdarcy",
        )

        ok = await module.async_migrate_entry(hass, entry)

        assert ok is True
        # ADB options carried over byte-compatibly
        assert entry.options["adb_enabled"] is True
        assert entry.options["adb_port"] == 5555
        assert entry.options["jellyfin_entity"] == "jf.x"
        # the v2 default, then the v3 address collapse
        assert entry.options[CONF_PROTOCOL_ENABLED] is False
        assert entry.options[CONF_HOST] == "10.0.0.31"
        assert "adb_host" not in entry.options
        assert entry.version == 3

    async def test_v2_prefers_the_protocol_host_over_the_adb_host(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(
            version=2,
            data={},
            options={
                CONF_PROTOCOL_ENABLED: True,
                CONF_PROTOCOL_HOST: "10.0.0.2",  # the address that was paired
                CONF_ADB_HOST: "10.0.0.1",
            },
            title="mdarcy",
        )

        assert await module.async_migrate_entry(hass, entry) is True

        assert entry.options[CONF_HOST] == "10.0.0.2"
        assert CONF_PROTOCOL_HOST not in entry.options
        assert CONF_ADB_HOST not in entry.options
        assert entry.options[CONF_PROTOCOL_ENABLED] is True  # untouched
        assert entry.version == 3

    async def test_v2_falls_back_to_the_adb_host(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(
            version=2,
            data={},
            options={CONF_PROTOCOL_ENABLED: False, CONF_ADB_HOST: "10.0.0.1"},
            title="mdarcy",
        )

        assert await module.async_migrate_entry(hass, entry) is True

        assert entry.options[CONF_HOST] == "10.0.0.1"
        assert CONF_ADB_HOST not in entry.options

    async def test_v2_without_an_address_migrates_to_an_empty_one(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(
            version=2,
            data={},
            options={CONF_PROTOCOL_ENABLED: False, CONF_ADB_HOST: "", CONF_PROTOCOL_HOST: ""},
            title="mdarcy",
        )

        assert await module.async_migrate_entry(hass, entry) is True

        assert entry.options[CONF_HOST] == ""
        assert CONF_ADB_HOST not in entry.options
        assert CONF_PROTOCOL_HOST not in entry.options
        assert entry.version == 3

    async def test_the_current_version_is_a_noop(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(
            version=3,
            data={},
            options={CONF_PROTOCOL_ENABLED: True, CONF_HOST: "shield.local"},
            title="mdarcy",
        )

        assert await module.async_migrate_entry(hass, entry) is True
        assert hass.config_entries.updating_entry_calls == []

    async def test_downgrade_rejected(self, tmp_path):
        module = _entry_module()
        hass = FakeHass(str(tmp_path))
        entry = SimpleNamespace(version=4, options={}, title="x")
        assert await module.async_migrate_entry(hass, entry) is False