managed-containerd/tests/test-runtime-state.py

#!/usr/bin/env python3

import ast
import json
import os
import sys
import tempfile
from pathlib import Path
from unittest.mock import patch

sys.path.insert(0, "daemon")

from registry import Plugin, Registry, RegistryError, Resource


def require(condition, message):
    if not condition:
        raise SystemExit("FAIL: " + message)


def write_plugin(
    root,
    name,
    manifest,
    source,
):
    plugin_root = Path(root) / name
    plugin_root.mkdir()

    (plugin_root / "manifest.json").write_text(
        json.dumps(manifest)
    )
    (plugin_root / "plugin.py").write_text(source)


print("===== GENERIC PLUGIN RUNTIME STATE =====")

with tempfile.TemporaryDirectory() as temporary_root:
    stateful_manifest = {
        "name": "stateful",
        "version": 1,
        "module": "plugin.py",
        "actions": ["monitor"],
        "required_config": ["REQUIRED_VALUE"],
        "optional_config": ["OPTIONAL_VALUE"],
    }

    write_plugin(
        temporary_root,
        "stateful",
        stateful_manifest,
        """
def monitor(config, runtime_state):
    counter = runtime_state.get("counter", 0) + 1
    runtime_state["counter"] = counter

    return {
        "counter": counter,
        "required": config["REQUIRED_VALUE"],
        "optional": config["OPTIONAL_VALUE"],
    }
""",
    )

    environment = {
        "MANAGED_CONTAINERD_PLUGIN_ROOT": (
            temporary_root
        ),
        "RESOURCE_ID": "stateful-resource",
        "PLUGIN": "stateful",
        "REQUIRED_VALUE": "required",
        "OPTIONAL_VALUE": "optional",
    }

    with patch.dict(
        os.environ,
        environment,
        clear=False,
    ):
        registry = Registry()

    first = registry.execute(
        "stateful-resource",
        "monitor",
    )
    second = registry.execute(
        "stateful-resource",
        "monitor",
    )

    require(
        first["result"]["counter"] == 1,
        "first stateful invocation has wrong counter",
    )
    require(
        second["result"]["counter"] == 2,
        "runtime state did not persist",
    )
    require(
        second["result"]["required"] == "required",
        "required configuration was not injected",
    )
    require(
        second["result"]["optional"] == "optional",
        "optional configuration was not injected",
    )

    resource = registry.resources[
        "stateful-resource"
    ]

    require(
        resource.runtime_state == {"counter": 2},
        "resource runtime state has unexpected contents",
    )

    print("stateful_handler       persisted")
    print("optional_config        injected")

    legacy_manifest = {
        "name": "legacy",
        "version": 1,
        "module": "plugin.py",
        "actions": ["monitor"],
        "required_config": ["REQUIRED_VALUE"],
    }

    write_plugin(
        temporary_root,
        "legacy",
        legacy_manifest,
        """
def monitor(config):
    return {
        "legacy": True,
        "required": config["REQUIRED_VALUE"],
    }
""",
    )

    legacy_plugin = Plugin(
        temporary_root,
        "legacy",
    )

    legacy_result = legacy_plugin.execute(
        "monitor",
        {"REQUIRED_VALUE": "required"},
        {},
    )

    require(
        legacy_result == {
            "legacy": True,
            "required": "required",
        },
        "one-argument plugin compatibility failed",
    )

    print("legacy_handler         compatible")

    try:
        Resource(
            "invalid-resource",
            legacy_plugin,
            {"REQUIRED_VALUE": None},
        )
    except RegistryError:
        print("required_config        enforced")
    else:
        raise SystemExit(
            "FAIL: missing required configuration accepted"
        )

production_manifest = json.loads(
    Path(
        "plugins/postgresql/manifest.json"
    ).read_text()
)

require(
    production_manifest.get("version") == 4,
    "PostgreSQL manifest version is not 4",
)
require(
    "PG_REPLICATION_PEERS"
    in production_manifest.get(
        "required_config",
        [],
    ),
    "replication peers are not required",
)
require(
    "PG_RECONNECT_GRACE_SECONDS"
    in production_manifest.get(
        "optional_config",
        [],
    ),
    "reconnect grace is not optional",
)

plugin_tree = ast.parse(
    Path(
        "plugins/postgresql/plugin.py"
    ).read_text()
)

monitor_functions = [
    node
    for node in plugin_tree.body
    if isinstance(node, ast.FunctionDef)
    and node.name == "monitor"
]

require(
    len(monitor_functions) == 1,
    "PostgreSQL monitor entry point is ambiguous",
)

monitor_arguments = [
    argument.arg
    for argument
    in monitor_functions[0].args.args
]

require(
    monitor_arguments[:2]
    == ["config", "runtime_state"],
    "PostgreSQL monitor does not accept runtime state",
)

print("postgresql_contract    validated")
print()
print("RUNTIME_STATE_TEST_PASS=1")