argus/src/agent/tests/test_config_metadata.py

from __future__ import annotations

import os
import unittest
from contextlib import contextmanager
from unittest.mock import patch

from app.config import AgentConfig, load_config


@contextmanager
def temp_env(**overrides: str | None):
    originals: dict[str, str | None] = {}
    try:
        for key, value in overrides.items():
            originals[key] = os.environ.get(key)
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        yield
    finally:
        for key, original in originals.items():
            if original is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = original


class LoadConfigMetadataTests(unittest.TestCase):
    @patch("app.config.Path.mkdir")
    def test_metadata_from_environment_variables(self, mock_mkdir):
        with temp_env(
            MASTER_ENDPOINT="http://master.local",
            AGENT_HOSTNAME="dev-user-one-pod",
            AGENT_ENV="prod",
            AGENT_USER="ops",
            AGENT_INSTANCE="node-1",
        ):
            config = load_config()

        self.assertEqual(config.environment, "prod")
        self.assertEqual(config.user, "ops")
        self.assertEqual(config.instance, "node-1")
        mock_mkdir.assert_called()

    @patch("app.config.Path.mkdir")
    def test_metadata_falls_back_to_hostname(self, mock_mkdir):
        with temp_env(
            MASTER_ENDPOINT="http://master.local",
            AGENT_HOSTNAME="qa-team-abc-pod-2",
            AGENT_ENV=None,
            AGENT_USER=None,
            AGENT_INSTANCE=None,
        ):
            config = load_config()

        self.assertEqual(config.environment, "qa")
        self.assertEqual(config.user, "team")
        self.assertEqual(config.instance, "abc")
        mock_mkdir.assert_called()

    @patch("app.config._load_metadata_from_state", return_value=("prod", "ops", "node-1"))
    @patch("app.config.Path.mkdir")
    def test_metadata_from_node_state(self, mock_mkdir, mock_state):
        with temp_env(
            MASTER_ENDPOINT="http://master.local",
            AGENT_HOSTNAME="host_abc",
            AGENT_ENV=None,
            AGENT_USER=None,
            AGENT_INSTANCE=None,
        ):
            config = load_config()

        self.assertEqual(config.environment, "prod")
        self.assertEqual(config.user, "ops")
        self.assertEqual(config.instance, "node-1")
        mock_state.assert_called_once()
        mock_mkdir.assert_called()

    @patch("app.config.Path.mkdir")
    def test_partial_environment_variables_fallback(self, mock_mkdir):
        with temp_env(
            MASTER_ENDPOINT="http://master.local",
            AGENT_HOSTNAME="stage-ml-001-node",
            AGENT_ENV="prod",
            AGENT_USER=None,
            AGENT_INSTANCE=None,
        ):
            config = load_config()

        self.assertEqual(config.environment, "stage")
        self.assertEqual(config.user, "ml")
        self.assertEqual(config.instance, "001")
        mock_mkdir.assert_called()

    @patch("app.config.Path.mkdir")
    def test_invalid_hostname_raises_error(self, mock_mkdir):
        with temp_env(
            MASTER_ENDPOINT="http://master.local",
            AGENT_HOSTNAME="invalidhostname",
            AGENT_ENV=None,
            AGENT_USER=None,
            AGENT_INSTANCE=None,
        ):
            with self.assertRaises(ValueError):
                load_config()

        mock_mkdir.assert_not_called()


class CollectMetadataTests(unittest.TestCase):
    @patch("app.collector._detect_ip_address", return_value="127.0.0.1")
    @patch("app.collector._detect_gpu_count", return_value=0)
    @patch("app.collector._detect_memory_bytes", return_value=1024)
    @patch("app.collector._detect_cpu_count", return_value=8)
    def test_collect_metadata_uses_config_fields(
        self,
        mock_cpu,
        mock_memory,
        mock_gpu,
        mock_ip,
    ):
        config = AgentConfig(
            hostname="dev-user-001-pod",
            environment="prod",
            user="ops",
            instance="node-1",
            node_file="/tmp/node.json",
            version="1.0.0",
            master_endpoint="http://master.local",
            report_interval_seconds=60,
            health_dir="/tmp/health",
        )

        from app.collector import collect_metadata

        metadata = collect_metadata(config)

        self.assertEqual(metadata["env"], "prod")
        self.assertEqual(metadata["user"], "ops")
        self.assertEqual(metadata["instance"], "node-1")
        self.assertEqual(metadata["hostname"], "dev-user-001-pod")
        self.assertEqual(metadata["ip"], "127.0.0.1")
        self.assertEqual(metadata["cpu_number"], 8)
        self.assertEqual(metadata["memory_in_bytes"], 1024)
        self.assertEqual(metadata["gpu_number"], 0)


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