Migrating the classes to pydantic

playbook/unittests/test_ActionRegistrar.py

@@ -0,0 +1,53 @@
+import os
+import unittest
+from unittest.mock import patch
+from playbook import ActionRegistrar, ActionRegistrarYAMLReader
+from playbook.action_registry import ActionRegistry
+
+
+CDIR = os.path.dirname(__file__)
+TEST_DATA = os.path.join(CDIR, "test_data")
+
+
+class TestActionRegistrar(unittest.TestCase):
+    @patch.object(ActionRegistrarYAMLReader, '_load_registries_from_file')
+    def test_reading_action_registry(self, mock_load):
+        mock_reg1 = ActionRegistry("docker_reg")
+        mock_reg2 = ActionRegistry("restic_reg")
+        mock_load.side_effect = [[mock_reg1], [mock_reg2]]
+
+        yaml_data = {
+            "registries": ["docker.py", "restic.py"]
+        }
+
+        expected_size = 2
+        expected_manifest = {
+            "registrar_manifest": [
+                {"registry_name": "docker_reg", "functions": {}},
+                {"registry_name": "restic_reg", "functions": {}}
+            ]
+        } 
+
+        reg: ActionRegistrar = ActionRegistrarYAMLReader.read(yaml_data)
+
+        self.assertEqual(expected_manifest, reg.manifest())
+        self.assertEqual(expected_size, reg.size())
+
+        pass
+
+    # def test_creating_action_registry(self):
+    #     yaml_data = {
+    #         "registries": ["docker.py", "restic.py"]
+    #     }
+    #
+    #     expected_size = 2
+    #     expected_manifest = {
+    #         "registrar_manifest": []
+    #     }
+    #
+    #     reg: ActionRegistrar = ActionRegistrarYAMLReader.read(
+    #         yaml_data, base_dir=os.path.join(TEST_DATA, "registries")
+    #     )
+    #
+    #     self.assertEqual(expected_manifest, reg.manifest())
+    #     self.assertEqual(expected_size, reg.size())

playbook/unittests/test_data/registries/docker.py

@@ -0,0 +1,65 @@
+from os import error
+
+from yaml import dump
+
+from playbook import models
+from playbook.action_registry import ActionRegistry
+from playbook.models import StepLog, ContextChecker
+
+import sys
+import json
+import docker
+import subprocess
+from typing import List, Dict, Any, Tuple
+# from docker.models.containers import Container
+
+
+dockeractions = ActionRegistry("docker_actions_reg")
+
+
+@dockeractions.register(name="get_service_container_name", version="")
+@ContextChecker.requires("service_name")
+def get_service_container_name(ctx, name) -> StepLog:
+    client = docker.from_env()
+
+    filters: Dict[str, Any]= {
+        "label": [f"com.docker.compose.service={ctx["service_name"]}"]
+    }
+    containers = client.containers.list(filters=filters, all=True)
+
+    if not containers:
+        return StepLog.fail(name, [{
+            "status": "failed", 
+            "output": f"no container found for service {ctx["service_name"]}"}
+        ])
+
+    new_data = {"container": containers[0].id}
+    return StepLog.ok(name, f"found container {containers[0].id}", pipe_ctx=new_data)
+
+
+@dockeractions.register(name="dump_container_pg_db", version="")
+@ContextChecker.requires("container", "dump_path", "db_user", "database")
+def dump_container_pg_db(ctx, name) -> StepLog:
+    client = docker.from_env()
+    container = client.containers.get(ctx["container"])
+
+    res = container.exec_run(
+        cmd=f"pg_dump -U {ctx["db_user"]} {ctx["database"]}",
+        stream=False,  # change to True to stream directly to host file (need to check)
+        demux=True,
+    )
+    stdout, stderr = res.output
+
+    if stderr:
+        if not isinstance(stderr, bytes):
+            return StepLog.fail(name, [{"status": "failed", "output": "stderr is not bytes"}])
+        if res.exit_code != 0:
+            error_msg = stderr.decode("utf-8") if stderr else "dump failed with no output"
+            return StepLog.fail(name, [{"status": "failed", "output": error_msg}])
+
+    if not isinstance(stdout, bytes):
+        return StepLog.fail(name, [{"status": "failed", "output": "stdout is not bytes"}])
+    with open(ctx["dump_path"], "wb") as f:
+        f.write(stdout)
+
+    return StepLog.ok(name, "database dump successful")    

playbook/unittests/test_data/registries/restic.py

@@ -0,0 +1,167 @@
+# restic actions registry
+from playbook.action_registry import ActionRegistry
+from playbook.models import ContextChecker, StepLog
+
+import sys
+import json
+import subprocess
+from typing import List, Dict, Any, Tuple
+
+
+# NOTE:
+# most restic run commands seem to be the same pattern of:
+# restic command -> parse output -> grab errors
+# we can create a class for this I reckon
+
+
+restic = ActionRegistry("restic_actions_reg")
+
+
+def parse_output(line: str) -> List[Dict[str, Any]]:
+    errors = []
+    for line in line.split('\n'):
+        if line.startswith("{"):  # } treesitter is borked lmao
+            data = json.loads(line)
+            errors.append(data)
+    return errors
+
+
+def find_restic_errors(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    errors = []
+    for s in data:
+        if s["message_type"] == "exit_error":
+            errors.append(s)
+    return errors
+
+
+def run_restic_command(cmd: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
+    """
+    Executes a restic command, streams output live to terminal, 
+    and returns (returncode, parsed_json_objects).
+    """
+    parsed_json = []
+    
+    # Start process with combined stdout/stderr pipe
+    process = subprocess.Popen(
+        cmd,
+        stdin=subprocess.DEVNULL,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1
+    )
+
+    # NOTE: is this what we want?
+    assert process.stdout is not None
+
+    # Read line-by-line in real-time
+    for line in iter(process.stdout.readline, ''):
+        # Write live output directly to terminal screen
+        sys.stdout.write(line)
+        sys.stdout.flush()
+
+        # Parse and capture valid JSON objects on the fly
+        clean_line = line.strip()
+        if clean_line.startswith('{'): # }
+            try:
+                data = json.loads(clean_line)
+                parsed_json.append(data)
+            except json.JSONDecodeError:
+                pass
+
+    process.stdout.close()
+    returncode = process.wait()
+    return returncode, parsed_json
+
+
+@restic.register(name="check_restic_environment", version="")
+def check_restic_environment(ctx, name: str = "") -> StepLog:
+    for v in ["passwordFile", "repoPath", "sourcePath", "serviceTag"]:
+        if not ctx.get(v):
+            return StepLog.fail(name, [f"{v} variable is empty"])
+
+    return StepLog.ok(name, "environment configuration seems to be valid")
+
+
+@restic.register(name="create_restic_repo", version="")
+def create_restic_repo(ctx, name: str = "") -> StepLog:
+    cmd = ["restic", "init", "--json", "-r", ctx["repoPath"], 
+         "--password-file", ctx["passwordFile"]]
+
+    returncode, json_output = run_restic_command(cmd)
+
+    errors: List[Dict[str, Any]] = find_restic_errors(json_output)
+
+    # Handle error messages in JSON output
+    if len(errors) > 0:
+        return StepLog.fail(name, errors)
+
+    # Make sure the return code is good too
+    if returncode != 0:
+        return StepLog.fail(name, [{"status": "failed", "output": json_output}])
+
+    return StepLog.ok(name, "repo create succesfully")
+
+
+@restic.register(name="create_repo_if_not_exists", version="")
+def create_repo_if_not_exists(ctx, name: str = "") -> StepLog:
+    log: StepLog = check_if_repo_exists(ctx, "")
+
+    msg: str = ""
+    if log.failed:
+        # Repo doesn't exist is error code 10
+        if log.error[0]["code"] != 10:
+            return StepLog.fail(name, [{"status": "failed", "output": log.error}])
+
+        # Create repo
+        repoCreateLog: StepLog = create_restic_repo(ctx, name + f".{create_restic_repo.__name__}")
+        if repoCreateLog.failed:
+            return StepLog.fail(name, [{"status": "failed", "output": repoCreateLog.error}])
+        msg = "repo created"
+    else:
+        msg = "repo already exists"
+
+
+    return StepLog.ok(name, msg)
+
+
+@restic.register(name="check_if_repo_exists", version="")
+def check_if_repo_exists(ctx, name: str = "") -> StepLog:
+    cmd = ["restic", "-r", ctx["repoPath"], "cat", "config", "--json",
+           "--password-file", ctx["passwordFile"]]
+    returncode, json_output = run_restic_command(cmd)
+
+    errors: List[Dict[str, Any]] = find_restic_errors(json_output)
+
+    # Handle error messages in JSON output
+    if len(errors) > 0:
+        return StepLog.fail(name, errors)
+
+    # Make sure the return code is good too
+    if returncode != 0:
+        return StepLog.fail(name, [{"status": "failed", "output": json_output}])
+
+    return StepLog.ok(name, "repo seems to exist")
+
+
+@restic.register(name="backup_data_to_restic_repo", version="")
+@ContextChecker.requires("source_path", "passwordFile")
+def backup_data_to_restic_repo(ctx, name) -> StepLog:
+    cmd = ["restic", "backup", ctx["source_path"], "--json", "--quiet",
+         "-r", ctx["repoPath"], "--password-file", ctx["passwordFile"]]
+    if ctx["tags"]:
+        cmd.extend(["--tag", ','.join(ctx["tags"])])
+
+    returncode, json_output = run_restic_command(cmd)
+    print(json_output, file=sys.stderr)
+    errors: List[Dict[str, Any]] = find_restic_errors(json_output)
+
+    # Handle error messages in JSON output
+    if len(errors) > 0:
+        return StepLog.fail(name, errors)
+
+    # Make sure the return code is good too
+    if returncode != 0:
+        return StepLog.fail(name, [{"status": "failed", "output": errors}])
+
+    return StepLog.ok(name, "backup successful")