Prepare once, reuse safely

import asyncio
from uuid import UUID

from contree_sdk import Contree


def require_success(run, label: str) -> None:
    if run.exit_code != 0:
        raise RuntimeError(f"{label} failed: {run.stderr}")


async def main() -> None:
    client = Contree()
    base = await client.images.oci("docker.io/library/alpine:3.19")
    if not isinstance(base.uuid, UUID):
        raise RuntimeError('Image import did not return a UUID')

    prepared = await base.run(
        shell="mkdir -p /workspace && printf 'prepared-v1\\n' > /workspace/prepared.txt",
        disposable=False,
        timeout=60,
    )
    require_success(prepared, "prepare")
    if prepared.uuid is None:
        raise RuntimeError("Preparation returned no reusable image")
    checkpoint_uuid = prepared.uuid

    checkpoint = await client.images.use(checkpoint_uuid, strict=True)
    first, second = await asyncio.gather(
        checkpoint.run(
            shell="test \"$(cat /workspace/prepared.txt)\" = prepared-v1 && printf 'A\\n'",
            disposable=True,
            timeout=30,
        ),
        checkpoint.run(
            shell="test \"$(cat /workspace/prepared.txt)\" = prepared-v1 && printf 'B\\n'",
            disposable=True,
            timeout=30,
        ),
    )
    require_success(first, "task A")
    require_success(second, "task B")
    if first.stdout != "A\n":
        raise RuntimeError('Task A did not return its expected output')
    if second.stdout != "B\n":
        raise RuntimeError('Task B did not return its expected output')
    if first.uuid is not None or second.uuid is not None:
        raise RuntimeError('Disposable tasks unexpectedly returned saved images')
    print(checkpoint_uuid)
    print(first.stdout + second.stdout, end="")


if __name__ == "__main__":
    asyncio.run(main())