herdr-session-switch/tests/test_discovery.py

from __future__ import annotations

import json
import os
from pathlib import Path
import subprocess
import tempfile
import unittest

from herdr_hub.cli import (
    DEFAULT_SESSION_NAME,
    Endpoint,
    decode_cmdline,
    dedupe_endpoints,
    deserialize_endpoint,
    discover_local_sessions,
    discover_remote_processes,
    parse_remote_argv,
    serialize_endpoint,
)


class EndpointTests(unittest.TestCase):
    def test_local_and_remote_commands(self) -> None:
        self.assertEqual(Endpoint.local().command("herdr"), ["herdr"])
        self.assertEqual(
            Endpoint.local("work").command("/bin/herdr"),
            ["/bin/herdr", "--session", "work"],
        )
        self.assertEqual(
            Endpoint.remote("orangepi", "work", "server").command("herdr"),
            [
                "herdr",
                "--remote",
                "orangepi",
                "--remote-keybindings",
                "server",
                "--session",
                "work",
            ],
        )
        self.assertEqual(
            Endpoint.remote(
                "ssh://me@windows:2222",
                "work",
                remote_transport="ssh",
            ).command("/bin/herdr", "/bin/ssh"),
            [
                "/bin/ssh",
                "-tt",
                "-o",
                "ServerAliveInterval=15",
                "-o",
                "ServerAliveCountMax=3",
                "-p",
                "2222",
                "-l",
                "me",
                "windows",
                "herdr --session work",
            ],
        )

    def test_identity_ignores_remote_keybinding_mode(self) -> None:
        local = Endpoint.remote("orangepi", remote_keybindings="local")
        server = Endpoint.remote("orangepi", remote_keybindings="server")
        self.assertEqual(local.key, server.key)
        self.assertEqual(dedupe_endpoints([local, server]), [local])
        self.assertEqual(
            Endpoint.remote("orangepi").key,
            Endpoint.remote("orangepi", remote_transport="ssh").key,
        )

    def test_short_labels_put_the_host_first_and_omit_default(self) -> None:
        self.assertEqual(Endpoint.local().short_label, "local")
        self.assertEqual(Endpoint.local("work").short_label, "local/work")
        self.assertEqual(Endpoint.remote("orangepi").short_label, "orangepi")
        self.assertEqual(Endpoint.remote("orangepi", "work").short_label, "orangepi/work")

    def test_default_local_sorts_first(self) -> None:
        endpoints = dedupe_endpoints(
            [
                Endpoint.remote("z"),
                Endpoint.local("work"),
                Endpoint.local(),
                Endpoint.remote("a"),
            ]
        )
        self.assertEqual(
            [endpoint.key for endpoint in endpoints],
            ["local:default", "local:work", "remote:a:default", "remote:z:default"],
        )

    def test_endpoint_metadata_and_legacy_command_round_trip(self) -> None:
        endpoint = Endpoint.remote(
            "ssh://me@host:2222",
            "work",
            "server",
            "ssh",
        )
        self.assertEqual(deserialize_endpoint(serialize_endpoint(endpoint)), endpoint)
        self.assertEqual(
            deserialize_endpoint(
                "",
                "exec /usr/local/bin/herdr --remote orangepi --session work",
            ),
            Endpoint.remote("orangepi", "work"),
        )


class RemoteArgvTests(unittest.TestCase):
    def test_separate_arguments(self) -> None:
        endpoint = parse_remote_argv(
            [
                "/home/me/.local/bin/herdr",
                "--remote",
                "orangepi",
                "--session",
                "work",
                "--remote-keybindings",
                "server",
            ]
        )
        self.assertEqual(endpoint, Endpoint.remote("orangepi", "work", "server"))

    def test_equal_arguments_and_arbitrary_order(self) -> None:
        endpoint = parse_remote_argv(
            ["herdr", "--session=work", "--remote=ssh://me@host:2222"]
        )
        self.assertEqual(endpoint, Endpoint.remote("ssh://me@host:2222", "work"))

    def test_default_session(self) -> None:
        endpoint = parse_remote_argv(["herdr", "--remote", "orangepi"])
        self.assertIsNotNone(endpoint)
        assert endpoint is not None
        self.assertEqual(endpoint.session, DEFAULT_SESSION_NAME)

    def test_non_herdr_and_malformed_commands_are_ignored(self) -> None:
        self.assertIsNone(parse_remote_argv(["bash", "-c", "herdr --remote host"]))
        self.assertIsNone(parse_remote_argv(["herdr", "--remote"]))
        self.assertIsNone(parse_remote_argv(["herdr", "--remote", "-bad"]))

    def test_cmdline_is_nul_delimited(self) -> None:
        self.assertEqual(
            decode_cmdline(b"/bin/herdr\0--remote\0host name\0"),
            ["/bin/herdr", "--remote", "host name"],
        )


class DiscoveryTests(unittest.TestCase):
    def test_local_discovery_returns_only_running_sessions(self) -> None:
        payload = {
            "sessions": [
                {"name": "default", "running": True},
                {"name": "work", "running": True},
                {"name": "old", "running": False},
            ]
        }

        def fake_run(*args, **kwargs):  # type: ignore[no-untyped-def]
            return subprocess.CompletedProcess(args[0], 0, json.dumps(payload), "")

        report = discover_local_sessions("herdr", run=fake_run)
        self.assertEqual(
            [endpoint.key for endpoint in report.endpoints],
            ["local:default", "local:work"],
        )
        self.assertEqual(report.warnings, [])

    def test_invalid_local_json_is_a_warning(self) -> None:
        def fake_run(*args, **kwargs):  # type: ignore[no-untyped-def]
            return subprocess.CompletedProcess(args[0], 0, "not json", "")

        report = discover_local_sessions("herdr", run=fake_run)
        self.assertEqual(report.endpoints, [])
        self.assertEqual(len(report.warnings), 1)

    def test_proc_discovery_reads_exact_argv_and_deduplicates(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            commands = {
                "101": ["/usr/bin/herdr", "--remote", "orangepi"],
                "102": ["herdr", "--remote=orangepi"],
                "103": ["herdr", "--remote", "prod", "--session", "work"],
                "104": ["bash", "-c", "herdr --remote ignored"],
            }
            for process_id, argv in commands.items():
                process = root / process_id
                process.mkdir()
                (process / "cmdline").write_bytes(
                    b"\0".join(os.fsencode(argument) for argument in argv) + b"\0"
                )

            report = discover_remote_processes(root, user_id=os.getuid())

        self.assertEqual(
            [endpoint.key for endpoint in report.endpoints],
            ["remote:orangepi:default", "remote:prod:work"],
        )


if __name__ == "__main__":
    unittest.main()