computer-playbook/tests/unit/python/plugins/lookup/test_compose_replicas.py

"""Unit tests for the compose_replicas lookup plugin.

Pins the SPOT contract for swarm replica calculation:

* swarm mode: emits `replicas: N` where N defaults to
  `len(groups[application_id])` or the override term.
* non-swarm mode: emits empty string.
* missing group / empty topology default: clamped to 1.
* an explicit override term is respected as-is, including 0 (lets a
  service stay unscheduled until the role scales it up).
"""

from __future__ import annotations

import importlib.util
import unittest

from ansible.errors import AnsibleError

from . import PROJECT_ROOT


def _load_lookup():
    spec = importlib.util.spec_from_file_location(
        "lookup_compose_replicas",
        str(PROJECT_ROOT / "plugins/lookup/compose_replicas.py"),
    )
    mod = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    spec.loader.exec_module(mod)
    return mod.LookupModule


class _DummyTemplar:
    def __init__(self, available_variables=None):
        self.available_variables = available_variables or {}

    def template(self, value):
        return value


class TestComposeReplicasLookup(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.LookupModule = _load_lookup()

    def _make(self, variables):
        lm = self.LookupModule()
        lm._templar = _DummyTemplar(variables)
        lm._loader = None
        return lm

    def test_compose_mode_no_arg_emits_empty(self):
        vars_ = {
            "DEPLOYMENT_MODE": "compose",
            "IS_COMPOSE_MODE": True,
            "IS_SWARM_MODE": False,
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["h1", "h2", "h3"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), [""])

    def test_compose_mode_with_arg_emits_empty(self):
        vars_ = {"DEPLOYMENT_MODE": "compose"}
        lm = self._make(vars_)
        self.assertEqual(lm.run([3], variables=vars_), [""])

    def test_missing_deployment_mode_defaults_to_compose(self):
        vars_ = {
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["h1", "h2"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), [""])

    def test_swarm_mode_default_uses_group_length(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["mgr-01", "wrk-01", "wrk-02"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), ["replicas: 3"])

    def test_swarm_mode_single_host_yields_one(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["h1"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), ["replicas: 1"])

    def test_swarm_mode_empty_group_falls_back_to_one(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {"web-app-x": []},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), ["replicas: 1"])

    def test_swarm_mode_missing_group_falls_back_to_one(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), ["replicas: 1"])

    def test_swarm_mode_explicit_override(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["h1", "h2", "h3"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([5], variables=vars_), ["replicas: 5"])

    def test_swarm_mode_override_string_coerced(self):
        vars_ = {"DEPLOYMENT_MODE": "swarm"}
        lm = self._make(vars_)
        self.assertEqual(lm.run(["2"], variables=vars_), ["replicas: 2"])

    def test_swarm_mode_override_zero_respected(self):
        vars_ = {"DEPLOYMENT_MODE": "swarm"}
        lm = self._make(vars_)
        self.assertEqual(lm.run([0], variables=vars_), ["replicas: 0"])

    def test_too_many_terms_raises(self):
        vars_ = {"DEPLOYMENT_MODE": "swarm"}
        lm = self._make(vars_)
        with self.assertRaises(AnsibleError):
            lm.run([1, 2], variables=vars_)

    def test_invalid_override_raises(self):
        vars_ = {"DEPLOYMENT_MODE": "swarm"}
        lm = self._make(vars_)
        with self.assertRaises(AnsibleError):
            lm.run(["not-a-number"], variables=vars_)

    def test_swarm_mode_whitespace_in_mode_normalised(self):
        vars_ = {
            "DEPLOYMENT_MODE": "  swarm  ",
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["h1", "h2"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run([], variables=vars_), ["replicas: 2"])

    def test_none_terms_treated_as_no_args(self):
        vars_ = {
            "DEPLOYMENT_MODE": "swarm",
            "IS_COMPOSE_MODE": False,
            "IS_SWARM_MODE": True,
            "application_id": "web-app-x",
            "groups": {"web-app-x": ["a", "b"]},
        }
        lm = self._make(vars_)
        self.assertEqual(lm.run(None, variables=vars_), ["replicas: 2"])


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