plugin_test.py

#!/usr/bin/env python3
import ast
import os
from functools import partial
from pathlib import Path
from typing import Optional

from test_base import TestBase

from warnet.k8s import download, wait_for_pod
from warnet.process import run_command


class PluginTest(TestBase):
    def __init__(self):
        super().__init__()
        self.network_dir = Path(os.path.dirname(__file__)) / "data" / "network_with_plugins"
        self.plugins_dir = Path(os.path.dirname(__file__)).parent / "resources" / "plugins"
        self.simln_exec = self.plugins_dir / "simln" / "plugin.py"

    def run_test(self):
        try:
            os.chdir(self.tmpdir)
            self.deploy_with_plugin()
            self.copy_results()
            self.assert_hello_plugin()
            self.test_capacity_multiplier_from_network_yaml()
        finally:
            self.cleanup()

    def deploy_with_plugin(self):
        self.log.info("Deploy the ln network with a SimLN plugin")
        results = self.warnet(f"deploy {self.network_dir}")
        self.log.info(results)
        wait_for_pod(self.get_first_simln_pod())

    def copy_results(self):
        pod = self.get_first_simln_pod()
        partial_func = partial(self.found_results_remotely, pod)
        self.wait_for_predicate(partial_func)

        download(pod, Path("/working/results"), Path("."))
        self.wait_for_predicate(self.found_results_locally)

    def found_results_remotely(self, pod: Optional[str] = None) -> bool:
        if pod is None:
            pod = self.get_first_simln_pod()
        self.log.info(f"Checking for results file in {pod}")
        results_file = run_command(f"{self.simln_exec} sh {pod} ls /working/results").strip()
        self.log.info(f"Results file: {results_file}")
        results = run_command(
            f"{self.simln_exec} sh {pod} cat /working/results/{results_file}"
        ).strip()
        self.log.info(results)
        return results.find("Success") > 0

    def get_first_simln_pod(self):
        command = f"{self.simln_exec} list-pod-names"
        pod_names_literal = run_command(command)
        self.log.info(f"{command}: {pod_names_literal}")
        pod_names = ast.literal_eval(pod_names_literal)
        return pod_names[0]

    def found_results_locally(self) -> bool:
        directory = "results"
        self.log.info(f"Searching {directory}")
        for root, _dirs, files in os.walk(Path(directory)):
            for file_name in files:
                file_path = os.path.join(root, file_name)

                with open(file_path) as file:
                    content = file.read()
                    if "Success" in content:
                        self.log.info(f"Found downloaded results in directory: {directory}.")
                        return True
        self.log.info(f"Did not find downloaded results in directory: {directory}.")
        return False

    def assert_hello_plugin(self):
        self.log.info("Waiting for the 'hello' plugin pods.")
        wait_for_pod("hello-pre-deploy")
        wait_for_pod("hello-post-deploy")
        wait_for_pod("hello-pre-network")
        wait_for_pod("hello-post-network")
        wait_for_pod("tank-0000-post-hello-pod")
        wait_for_pod("tank-0000-pre-hello-pod")
        wait_for_pod("tank-0001-post-hello-pod")
        wait_for_pod("tank-0001-pre-hello-pod")
        wait_for_pod("tank-0002-post-hello-pod")
        wait_for_pod("tank-0002-pre-hello-pod")
        wait_for_pod("tank-0003-post-hello-pod")
        wait_for_pod("tank-0003-pre-hello-pod")
        wait_for_pod("tank-0004-post-hello-pod")
        wait_for_pod("tank-0004-pre-hello-pod")
        wait_for_pod("tank-0005-post-hello-pod")
        wait_for_pod("tank-0005-pre-hello-pod")

    def test_capacity_multiplier_from_network_yaml(self):
        """Test that the capacity multiplier from network.yaml is properly applied."""
        self.log.info("Testing capacity multiplier from network.yaml configuration...")

        # Get the first simln pod
        pod = self.get_first_simln_pod()

        # Wait a bit for simln to start and generate activity
        import time

        time.sleep(10)

        # Check the sim.json file to verify the configuration is correct
        sim_json_content = run_command(f"{self.simln_exec} sh {pod} cat /working/sim.json")

        # Parse the JSON to check for capacityMultiplier
        import json

        try:
            sim_config = json.loads(sim_json_content)
            if "capacityMultiplier" not in sim_config:
                self.fail("capacityMultiplier not found in sim.json configuration")

            expected_multiplier = 5  # As configured in network.yaml
            if sim_config["capacityMultiplier"] != expected_multiplier:
                self.fail(
                    f"Expected capacityMultiplier {expected_multiplier}, got {sim_config['capacityMultiplier']}"
                )

            self.log.info(
                f"✓ Found capacityMultiplier {sim_config['capacityMultiplier']} in sim.json"
            )

        except json.JSONDecodeError as e:
            self.fail(f"Invalid JSON in sim.json: {e}")

        # Try to get logs from the simln container (but don't fail if it hangs)
        logs = ""
        try:
            # Use kubectl logs (more reliable)
            logs = run_command(f"kubectl logs {pod} --tail=50")
        except Exception as e:
            self.log.warning(f"Could not get logs from simln container: {e}")
            self.log.info(
                "✓ Simln container is running with correct capacityMultiplier configuration"
            )
            self.log.info(
                "✓ Skipping log analysis due to log access issues, but configuration is correct"
            )
            return

        # Look for multiplier information in the logs
        if "multiplier" not in logs:
            self.log.warning(
                "No multiplier information found in simln logs, but this might be due to timing"
            )
            self.log.info(
                "✓ Simln container is running with correct capacityMultiplier configuration"
            )
            return

        # Check that we see the expected multiplier value (5 as configured in network.yaml)
        if "with multiplier 5" not in logs:
            self.log.warning(
                "Expected multiplier value 5 not found in simln logs, but this might be due to timing"
            )
            self.log.info(
                "✓ Simln container is running with correct capacityMultiplier configuration"
            )
            return

        # Verify that activity is being generated (should see "payments per month" or "payments per hour")
        if "payments per month" not in logs and "payments per hour" not in logs:
            self.log.warning(
                "No payment activity generation found in simln logs, but this might be due to timing"
            )
            self.log.info(
                "✓ Simln container is running with correct capacityMultiplier configuration"
            )
            return

        self.log.info("✓ Capacity multiplier from network.yaml is being applied correctly")
        self.log.info("Capacity multiplier from network.yaml test completed successfully")


if __name__ == "__main__":
    test = PluginTest()
    test.run_test()