pve_tests/main.py

#!/usr/bin/env python3

import argparse
import logging
import os
import time

import urllib3
import yaml
from dotenv import load_dotenv
from proxmoxer import ProxmoxAPI, ResourceException
from proxmoxer.tools import Tasks

from functions import clone_template, delete_vm, get_vm_ip

urllib3.disable_warnings()

FORMAT = "%(asctime)s %(name)s %(levelname)s %(message)s"
logging.basicConfig(format=FORMAT)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


load_dotenv()
PROXMOX_HOST: str | None = os.environ.get("PROXMOX_HOST")
PROXMOX_USER: str | None = os.environ.get("PROXMOX_USER")
PROXMOX_USER_FULL: str | None = os.environ.get("PROXMOX_USER_FULL")
PROXMOX_PASSWORD: str | None = os.environ.get("PROXMOX_PASSWORD")


def tasks_wait_completion(proxmox: ProxmoxAPI, upids: list[str]) -> None:
    for upid in upids:
        Tasks.blocking_status(proxmox, upid)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "branch",
        help="used to pick the appropriate VM template to clone",
        choices=["sisyphus", "c10f2", "p10"],
    )
    return parser.parse_args()


def main() -> None:
    assert PROXMOX_HOST is not None
    assert PROXMOX_USER is not None
    assert PROXMOX_USER_FULL is not None
    assert PROXMOX_PASSWORD is not None

    args = parse_args()

    proxmox = ProxmoxAPI(
        PROXMOX_HOST,
        user=PROXMOX_USER_FULL,
        password=PROXMOX_PASSWORD,
        verify_ssl=False,
    )

    nodes: list[dict[str, str | int]] = proxmox.nodes.get()  # pyright: ignore
    nodes = list(filter(lambda node: node["status"] == "online", nodes))
    # node = min(nodes, key=lambda node: node['disk'])['node']
    node: str = "pve05"

    match args.branch:
        case "sisyphus":
            template_id: int = 374
        case "p10":
            template_id: int = 520
        case "c10f2":
            template_id: int = 374
        case _:
            assert False

    prefix: str = PROXMOX_USER + "-test-k8s"
    vm_names: dict[str, str] = {
        "master": f"{prefix}-master",
        "worker1": f"{prefix}-node1",
        "worker2": f"{prefix}-node2",
    }
    # vm_ids = {
    # 'master': 500,
    # 'worker1': 501,
    # 'worker2': 502,
    # }
    vm_ids: dict[str, int] = {
        "master": 510,
        "worker1": 511,
        "worker2": 512,
    }

    delete: str = os.environ["DELETE_VMS"]
    assert delete in ("0", "1")

    if delete != "0":
        logger.info("Stopping VMs %s...", vm_ids)
        tasks_wait_completion(
            proxmox,
            [
                proxmox.nodes(node).qemu(vm_ids["master"]).status.stop.post(),
                proxmox.nodes(node).qemu(vm_ids["worker1"]).status.stop.post(),
                proxmox.nodes(node).qemu(vm_ids["worker2"]).status.stop.post(),
            ],
        )

        logger.info("Deleting VMs %s...", vm_ids)
        tasks_wait_completion(
            proxmox,
            [
                delete_vm(proxmox, node, vm_ids["master"]),
                delete_vm(proxmox, node, vm_ids["worker1"]),
                delete_vm(proxmox, node, vm_ids["worker2"]),
            ],
        )

    vm_id_master, upid_master = clone_template(
        proxmox,
        node,
        template_id,
        newid=vm_ids["master"],
        name=vm_names["master"],
    )
    assert vm_id_master == vm_ids["master"]

    vm_id_worker1, upid_worker1 = clone_template(
        proxmox,
        node,
        template_id,
        newid=vm_ids["worker1"],
        name=vm_names["worker1"],
    )
    assert vm_id_worker1 == vm_ids["worker1"]

    vm_id_worker2, upid_worker2 = clone_template(
        proxmox,
        node,
        template_id,
        newid=vm_ids["worker2"],
        name=vm_names["worker2"],
    )
    assert vm_id_worker2 == vm_ids["worker2"]

    logger.info("Waiting for cloning to complete...")
    tasks_wait_completion(
        proxmox,
        [
            upid_master,
            upid_worker1,
            upid_worker2,
        ],
    )
    logger.info("Cloning completed!")

    logger.info("Starting VMs...")
    tasks_wait_completion(
        proxmox,
        [
            proxmox.nodes(node).qemu(vm_ids["master"]).status.start.post(),
            proxmox.nodes(node).qemu(vm_ids["worker1"]).status.start.post(),
            proxmox.nodes(node).qemu(vm_ids["worker2"]).status.start.post(),
        ],
    )
    logger.info("VMs are running!")

    logger.info("Waiting for QEMU guest agent to start")
    done = False
    while not done:
        try:
            _ = proxmox.nodes(node).qemu(vm_ids["master"]).agent.ping.post()
            _ = proxmox.nodes(node).qemu(vm_ids["worker1"]).agent.ping.post()
            _ = proxmox.nodes(node).qemu(vm_ids["worker2"]).agent.ping.post()
        except ResourceException:
            time.sleep(1)
        except Exception as e:
            raise RuntimeError("Failed to ping QEMU agent") from e
        else:
            done = True
    logger.info("QEMU guest agent is running")

    vm_ip_master: str = get_vm_ip(proxmox, node, vm_ids["master"])
    vm_ip_worker1: str = get_vm_ip(proxmox, node, vm_ids["worker1"])
    vm_ip_worker2: str = get_vm_ip(proxmox, node, vm_ids["worker2"])
    proxmox.logout()

    tmp_path: str = "./tmp"
    if not os.path.exists(tmp_path):
        os.makedirs(tmp_path)

    with open(f"{tmp_path}/vm_ids", "w") as ofile:
        vm_ids_file = f'{vm_ids["master"]}\t{vm_names["master"]}\n'
        vm_ids_file += f'{vm_ids["worker1"]}\t{vm_names["worker1"]}\n'
        vm_ids_file += f'{vm_ids["worker2"]}\t{vm_names["worker2"]}\n'
        logger.info(vm_ids_file)
        ofile.write(vm_ids_file)

    with open(f"{tmp_path}/hosts", "w") as ofile:
        hosts = f'{vm_ip_master}\t{vm_names["master"]}\n'
        hosts += f'{vm_ip_worker1}\t{vm_names["worker1"]}\n'
        hosts += f'{vm_ip_worker2}\t{vm_names["worker2"]}\n'
        logger.info(hosts)
        ofile.write(hosts)

    inventory = {
        "workers": {
            "hosts": {
                "worker1": {
                    "ansible_host": vm_ip_worker1,
                },
                "worker2": {
                    "ansible_host": vm_ip_worker2,
                },
            },
        },
        "all_vms": {
            "hosts": {
                "master": {
                    "ansible_host": vm_ip_master,
                },
            },
            "children": {"workers": None},
            "vars": {"ansible_user": "root"},
        },
    }

    with open(f"{tmp_path}/generated_inventory.yaml", "w") as ofile:
        yaml.dump(inventory, ofile)


if __name__ == "__main__":
    main()