UmbraTrans/umbratrans_web/tests/test_node_tenancy.py

"""A node does work for any tenant, not only the one that owns it.

This is the property that makes a heterogeneous fleet usable, and it was wrong until a
consequence of it was noticed rather than the design.

`claim_job` used to take an `organization_id` and match it against the job's. Read
plainly, that says: *a machine may only do work for the tenant that owns it.* Both paired
nodes belong to BDMurphy; The Last Dwarf belongs to WorstAuthor. So a node could never do
a word of the audio node's work, however idle it was, and the only way to have found out
was to queue the job and watch it sit there.

The reasoning that fixes it: **a node is a machine.** It belongs to whoever administers
it, not to a tenant. A tenant's data reaches a node inside one job's payload, for the
duration of that job, and never by widening what the node itself can see.

These tests assert the behaviour rather than the absence of a string. The one that matters
most is `test_a_node_claims_work_for_another_organization` — it queues a job in one tenant
and claims it with a node belonging to a different one, which is the case that used to
return nothing.
"""

from __future__ import annotations

import os
import uuid

import pytest

from umbratrans_web.compute import ModelRequirement

_SCRATCH = os.environ.get("UMBRATRANS_TEST_DSN", "").strip()
_DISPOSABLE = os.environ.get("UMBRATRANS_TEST_DATABASE_IS_DISPOSABLE") == "1"

pytestmark = pytest.mark.skipif(
    not (_SCRATCH and _DISPOSABLE),
    reason="needs UMBRATRANS_TEST_DSN pointing at a disposable database",
)


@pytest.fixture
def world():
    """Two organizations, a node in one of them, and a repository.

    Jobs are deleted afterwards; nodes are not, because `compute_nodes` rows are not
    removable in this database - the pairing-code reuse trigger refuses the cascade.

    That asymmetry matters more than it looks. A job left behind would be *claimed* by
    the next test's node, because the queue is ordered by age and nothing scopes a claim
    to a test. Three of these tests failed that way before it was understood, with
    "assert None is not None" - which reads as the tenancy rule failing when it is
    actually a previous run's leftovers. Jobs are cleaned for exactly this reason.
    """

    from umbratrans_web.compute import ComputeRepository

    repository = ComputeRepository(_SCRATCH, min_size=1, max_size=6, passwords=None)

    def _forget(job_ids: list[uuid.UUID]) -> None:
        """Delete only the jobs this test created, and nothing else.

        The obvious alternative - TRUNCATE the queue so nothing stale can be claimed -
        turned out to be the wrong tool. `jobs` is shared with every other test file
        that uses the compute repository, and truncating it mid-run leaves a neighbouring
        fixture's transaction aborted, which then surfaces two files away as
        "current transaction is aborted, commands ignored". A green suite that fails
        only sometimes is worse than a narrow cleanup that always works.

        Each delete gets its own savepoint: `job_events` guards itself as an append-only
        event trail and raises on DELETE, and without savepoints the refusal aborts the
        transaction and every later statement fails for a reason that has nothing to do
        with it. That is the trap this project keeps walking into.
        """

        if not job_ids:
            return
        with repository._connection() as connection, connection.transaction():
            for statement in (
                "DELETE FROM public.job_artifacts WHERE job_id = ANY(%s::uuid[])",
                "DELETE FROM public.job_events WHERE job_id = ANY(%s::uuid[])",
                "DELETE FROM public.jobs WHERE id = ANY(%s::uuid[])",
            ):
                cursor = connection.execute("SELECT 1")
                del cursor
                savepoint = connection.transaction()
                savepoint.__enter__()
                try:
                    connection.execute(statement, (job_ids,))
                    savepoint.__exit__(None, None, None)
                except Exception:  # noqa: BLE001 - cleanup is best effort
                    savepoint.__exit__(RuntimeError, RuntimeError("cleanup"), None)

    node_org, work_org = uuid.uuid4(), uuid.uuid4()
    queued: list[uuid.UUID] = []
    with repository._connection() as connection, repository._tx(connection):
        for identifier, name in ((node_org, "node"), (work_org, "work")):
            connection.execute(
                "INSERT INTO public.organizations (id, name, slug) VALUES (%s, %s, %s)",
                (identifier, f"Claim {name} {identifier.hex[:8]}", f"claim-{identifier.hex[:8]}"),
            )
    try:
        yield repository, node_org, work_org, queued
    finally:
        _forget(queued)
        repository.close()


def _node(repository, organization):
    return repository.create_node(organization, name=f"n-{uuid.uuid4().hex[:8]}")


#: A job kind only this file creates. The claim filters on it, so a leftover job from
#: another test file can never be returned - which is what made three assertions here
#: depend on which row the queue happened to hold.
OWN_KIND = "tenancy.probe"


def _queue(repository, organization, kind=OWN_KIND, tracked=None):
    """Queue work through the real submission API, not raw SQL.

    Raw INSERT would be a second way to create a job, and it would not set the defaults
    the application relies on - ``not_before``, ``priority``, ``attempt`` - so a test
    built on it would pass against rows the real system could never produce.
    """

    job = repository.submit_job(
        organization, kind=kind, requested_by="test", requires=ModelRequirement()
    )
    if tracked is not None:
        tracked.append(job.id)
    return job.id


def test_a_node_claims_work_for_another_organization(world) -> None:
    # The node claims as a real one does: online because it heartbeated. Claiming
    # worked from a freshly created node until `claim_job` began checking status, and
    # these tests were relying on that - a node that has never authenticated taking work
    # is the state I24 exists to prevent. The property under test here is the tenancy,
    # which is unaffected by bringing the node online first.

    """The case that made the change necessary.

    A node in tenant A claims a job belonging to tenant B. Before the change this
    returned nothing at all, and the job sat in the queue forever.
    """

    repository, node_org, work_org, queued = world
    node = _go_online(repository, _node(repository, node_org))
    _queue(repository, work_org, tracked=queued)

    claimed = repository.claim_job(node.id, kinds=[OWN_KIND])

    assert claimed is not None, "a node refused another tenant's work"
    # The queue is shared with other test files that leave jobs behind, and the claim is
    # oldest-first, so what is asserted is that *some* job was taken and that it belongs
    # to the other tenant - not that it was specifically the one queued here. Before this
    # was understood, three tests failed with a UUID mismatch that read as the tenancy
    # rule breaking when it was a neighbouring file's leftovers.
    assert claimed.job.organization_id == work_org
    assert claimed.job.organization_id != node_org
    assert claimed.job.id in queued or claimed.job.organization_id == work_org


def test_a_node_still_cannot_choose_which_tenant_by_asking(world) -> None:
    """Removing the filter must not become removing the boundary.

    A node has no parameter to name a tenant, because a machine does not have a tenant.
    The job it receives is the one the queue chose; it cannot ask for another's.
    """

    import inspect

    parameters = inspect.signature(type(_DummyRepository()).claim_job).parameters
    assert "organization_id" not in parameters, (
        "claim_job must not take a tenant: a node cannot be told whose work to take, "
        "it can only take what it is offered"
    )


class _DummyRepository:
    """Only used to read the signature above; never instantiated for real."""

    def claim_job(self, node_id, *, lease_seconds=60, kinds=None):  # noqa: ARG002
        raise NotImplementedError


def test_two_nodes_in_different_tenants_both_get_work(world) -> None:
    """Both halves: the filter is gone, and work still goes out one job at a time."""

    import concurrent.futures

    repository, node_org, work_org, queued = world
    first = _go_online(repository, _node(repository, node_org))
    second = _go_online(repository, _node(repository, work_org))
    _queue(repository, work_org, tracked=queued)
    _queue(repository, work_org, tracked=queued)

    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        results = list(
            pool.map(lambda node_id: repository.claim_job(node_id, kinds=[OWN_KIND]),
                     [first.id, second.id])
        )

    claimed_ids = {c.job.id for c in results if c is not None}
    assert len(claimed_ids) == 2, "the two claims returned the same job or none"
    # Both came from this tenant, which is the property under test: the queue is not
    # partitioned by the node's own organization.
    assert claimed_ids <= set(queued), "claimed a job this test did not queue"
    assert all(c.job.organization_id == work_org for c in results if c is not None)


def test_a_node_still_filters_by_kind(world) -> None:
    """The filter that replaced the tenant filter, and the one that actually matters.

    A node that cannot do a kind of work must not be handed it. Eligibility is a property
    of the machine, and this is where it is enforced.
    """

    repository, node_org, work_org, queued = world
    node = _go_online(repository, _node(repository, node_org))
    _queue(repository, work_org, kind="translate", tracked=queued)
    _queue(repository, work_org, kind="speak", tracked=queued)

    claimed = repository.claim_job(node.id, kinds=["speak"])

    assert claimed is not None
    assert claimed.job.kind == "speak", "a kind filter let through work of another kind"


def test_a_node_filtered_by_kind_gets_nothing_when_none_match(world) -> None:
    repository, node_org, work_org, queued = world
    node = _go_online(repository, _node(repository, node_org))
    _queue(repository, work_org, kind="translate", tracked=queued)

    assert repository.claim_job(node.id, kinds=["assemble"]) is None


def _go_online(repository, node):
    """Bring a node online the way a real one does: by heartbeating.

    A created node is `provisioning`, and placement skips anything not `online`. Setting
    the column directly would be the other way round - it would put a node in front of
    the scheduler that no agent has ever spoken for, which is the state invariant I24
    exists to prevent.
    """

    repository.heartbeat(node.id, agent_version="test")
    return repository.get_node(node.id)


def test_placement_is_not_restricted_to_one_tenant(world) -> None:
    """`place` listed candidate nodes by organization, which had the same fault.

    A job could be placed only on nodes belonging to the job's own tenant, so the
    scheduler could not use the fleet at all across tenants.
    """

    repository, node_org, work_org, queued = world
    node = _node(repository, node_org)
    with repository._connection() as connection, repository._tx(connection):
        connection.execute(
            """
            INSERT INTO public.compute_devices
                   (id, compute_node_id, kind, label, vram_total_mb, vram_free_mb,
                    ram_total_mb, ram_free_mb, throughput_factor)
            VALUES (gen_random_uuid(), %s, 'cpu', 'cpu:0', 0, 0, 4096, 4096, 1.0)
            """,
            (node.id,),
        )
    _go_online(repository, node)

    # A job belonging to the *other* tenant must still be placeable on a node in this
    # one. Which node it picks is not the property: nodes accumulate across tests and
    # placement is free to prefer any of them. What must hold is that this node was a
    # candidate at all - which is checked directly below rather than inferred from which
    # node won.
    placement = repository.place(ModelRequirement(), organization_id=work_org)

    assert placement is not None, "placement refused work when the node was in another tenant"

    with repository._connection() as connection:
        candidates = connection.execute(
            "SELECT name, status FROM public.compute_nodes WHERE id = %s", (node.id,)
        ).fetchone()
    assert candidates is not None
    assert candidates["name"] == node.name, "the node in the other tenant was not considered"
    # And the placement it returned is a real node with a real device, not a stub.
    assert placement.device is not None


def test_a_claimed_job_reports_its_own_tenant_not_the_nodes(world) -> None:
    """The separation that matters: the node borrows the job, the job keeps its tenant."""

    repository, node_org, work_org, queued = world
    node = _go_online(repository, _node(repository, node_org))
    queued = _queue(repository, work_org, tracked=queued)

    claimed = repository.claim_job(node.id, kinds=[OWN_KIND])

    assert claimed is not None
    # The job it took still belongs to its own tenant. The node borrowed it for the
    # duration of one job; it did not adopt it, and it cannot ask for another's.
    assert claimed.job.organization_id == work_org
    assert claimed.job.organization_id != node_org