diff --git a/src/hosts/api.py b/src/hosts/api.py index fef7d79..42882fb 100644 --- a/src/hosts/api.py +++ b/src/hosts/api.py @@ -43,9 +43,7 @@ async def create_host( ] = None, ) -> Host: host_create = payload or HostCreate() - # An omitted expires_at gets the default lease; an explicit null is the - # caller's deliberate opt-in to a permanent host. - expires_at = host_create.expires_at if "expires_at" in host_create.model_fields_set else ... + permanent = "expires_at" in host_create.model_fields_set and not host_create.expires_at try: return await service.get_or_create_host( @@ -54,7 +52,8 @@ async def create_host( secrets={name: entry.to_storage() for name, entry in host_create.secrets.items()}, image=host_create.image, template=host_create.template, - expires_at=expires_at, + expires_at=host_create.expires_at, + permanent=permanent, idempotency_key=idempotency_key, provider=host_create.provider, instance_type=host_create.instance_type, diff --git a/src/hosts/service.py b/src/hosts/service.py index a2ab03d..18a458e 100644 --- a/src/hosts/service.py +++ b/src/hosts/service.py @@ -4,7 +4,6 @@ import time import uuid from datetime import UTC, datetime, timedelta -from types import EllipsisType from typing import Any import httpx @@ -101,17 +100,13 @@ async def get_or_create_host( secrets: dict[str, dict[str, Any]] | None = None, image: str | None, template: uuid.UUID | None = None, - expires_at: datetime | None | EllipsisType = ..., + expires_at: datetime | None = None, + permanent: bool = False, idempotency_key: str | None = None, provider: str | None = None, instance_type: str | None = None, disk_gb: int | None = None, ) -> Host: - # ``...`` (omitted) means "default lease"; an explicit None is the - # caller's deliberate opt-in to a permanent, never-reaped host. The - # sentinel travels to the point where a lease is actually stamped - # (pool claim, or the post-provision rewrite) so the in-flight row - # keeps the short provisioning safety TTL. if provider: registered = get_provider_names() if provider not in registered: @@ -130,7 +125,10 @@ async def get_or_create_host( customized = env or secrets or image or template or instance_type or disk_gb if not customized and self.settings.get_pool_targets().get(requested_provider): host = await self._try_claim_pool_host( - service_account=service_account, provider=requested_provider, expires_at=expires_at + service_account=service_account, + provider=requested_provider, + expires_at=expires_at, + permanent=permanent, ) if not host: host = await self.create_host( @@ -140,6 +138,7 @@ async def get_or_create_host( image=image, template=template, expires_at=expires_at, + permanent=permanent, provider=provider, instance_type=instance_type, disk_gb=disk_gb, @@ -164,7 +163,8 @@ async def _try_claim_pool_host( *, service_account: str | None = None, provider: str, - expires_at: datetime | None | EllipsisType, + expires_at: datetime | None = None, + permanent: bool = False, ) -> Host | None: # The claim's WHERE on claimed_at is the race guard: concurrent # claimants resolve to one winner per row on both dialects. @@ -180,15 +180,12 @@ async def _try_claim_pool_host( .order_by(Host.created_at.asc()) .limit(1) ) - if vm.max_lifetime and isinstance(expires_at, datetime): + if vm.max_lifetime and expires_at: # A warm host that the provider stops before the requested end is no match. candidates = candidates.where(Host.created_at >= expires_at - vm.max_lifetime) if candidate := (await self.session.execute(candidates)).scalar_one_or_none(): - if expires_at is ...: - expires_at = self._default_lease_expires_at() - # The claim replaces the warm-pool max-age TTL with the caller's lease: - # a concrete window (explicit or the default), or None for a caller - # who deliberately opted into a permanent host. + # The claim replaces the warm-pool max-age TTL with the caller's lease. + expires_at = None if permanent else expires_at or self._default_lease_expires_at() result = await self.session.execute( update(Host) .where(Host.id == candidate.id) @@ -215,7 +212,8 @@ async def create_host( secrets: dict[str, dict[str, Any]] | None = None, image: str | None, template: uuid.UUID | None = None, - expires_at: datetime | None | EllipsisType = ..., + expires_at: datetime | None = None, + permanent: bool = False, provider: str | None = None, instance_type: str | None = None, disk_gb: int | None = None, @@ -253,17 +251,10 @@ async def create_host( image = found.image uid = uuid7() now = utc_now() - # Safety TTL covers the strand window: if the client disconnects - # mid-provision, this is what makes the janitor reap the row + VM. - # Replaced with the caller's value after provisioning succeeds. A - # default-lease create keeps just the safety TTL in flight — the - # lease is stamped only once the host is usable. + # The safety TTL makes the janitor reap the row and VM if the client + # disconnects mid-provision; the lease is stamped once the host is usable. safety_expires_at = now + timedelta(seconds=self.settings.provisioning_grace_seconds) - initial_expires_at = ( - max(expires_at, safety_expires_at) - if isinstance(expires_at, datetime) - else safety_expires_at - ) + initial_expires_at = max(expires_at, safety_expires_at) if expires_at else safety_expires_at host = Host( id=uid, service_account=service_account, @@ -309,14 +300,10 @@ async def create_host( if host.status == HostStatus.ERROR.value: raise ProvisioningFailedError(host.last_error) - # Provisioning won: replace the safety TTL with the caller's intent - # in a dedicated session so we don't extend ``self.session``'s - # transaction (which can perturb advisory-lock-bearing callers like - # the pool maintainer). Guarded on the in-flight value: a renewal - # that landed while the host was bootstrapping is newer intent and - # must not be clobbered. - if expires_at is ...: - expires_at = self._default_lease_expires_at() + # A dedicated session, so self.session's transaction stays short for + # callers that hold an advisory lock. The match on the in-flight value + # spares a renewal that landed while the host was bootstrapping. + expires_at = None if permanent else expires_at or self._default_lease_expires_at() async with async_session_factory() as ttl_session: await ttl_session.execute( update(Host) diff --git a/src/hosts/tests/test_api.py b/src/hosts/tests/test_api.py index c7b33d4..4cc5530 100644 --- a/src/hosts/tests/test_api.py +++ b/src/hosts/tests/test_api.py @@ -88,23 +88,23 @@ async def test_create_host_creates_record_and_enqueues_provisioning(client, monk payload = response.json() assert uuid.UUID(payload["id"]).version == 7 assert payload["name"].startswith("sb-") - assert payload["status"] == "provisioning" # provision mock is a no-op + assert payload["status"] == "provisioning" assert payload["provider"] == "exe" assert payload["image"] == ExeSettings().default_image # pyright: ignore[reportCallIssue] # internal_ssh_host is populated by provision() (mocked out here), so # the response shows the row's pre-provision state. - assert payload["internal_ssh_host"] is None - assert payload["tailscale_device_id"] is None + assert not payload["internal_ssh_host"] + assert not payload["tailscale_device_id"] assert "env" not in payload async with async_session_factory() as session: host = await session.get(Host, uuid.UUID(payload["id"])) - assert host is not None + assert host assert host.image == ExeSettings().default_image # pyright: ignore[reportCallIssue] assert host.env == {} mocked_provision.assert_awaited_once() - assert mocked_provision.await_args is not None + assert mocked_provision.await_args assert mocked_provision.await_args.args[0].id == host.id @@ -131,11 +131,11 @@ async def test_create_host_persists_env_without_returning_it(client, monkeypatch async with async_session_factory() as session: host = await session.get(Host, uuid.UUID(payload["id"])) - assert host is not None + assert host assert host.env == env mocked_provision.assert_awaited_once() - assert mocked_provision.await_args is not None + assert mocked_provision.await_args assert mocked_provision.await_args.args[0].id == host.id @@ -156,7 +156,7 @@ async def test_create_host_stores_and_returns_custom_image(client, monkeypatch): async with async_session_factory() as session: host = await session.get(Host, uuid.UUID(payload["id"])) - assert host is not None + assert host assert host.image == image @@ -173,16 +173,15 @@ async def test_create_host_stores_and_returns_expires_at(client, monkeypatch): assert response.status_code == 201 payload = response.json() - assert payload["expires_at"] is not None + assert payload["expires_at"] async with async_session_factory() as session: host = await session.get(Host, uuid.UUID(payload["id"])) - assert host is not None - assert host.expires_at is not None + assert host + assert host.expires_at async def test_create_host_rejects_naive_expires_at(client): - # ISO 8601 without a tz offset is a naive datetime; the validator must reject. response = await client.post( "/hosts", headers={"Authorization": "Bearer service-token"}, @@ -241,7 +240,7 @@ async def capture_inflight_expiry(self, host): after = utc_now() assert before + grace <= inflight["expires_at"] <= after + grace - assert host.expires_at is not None + assert host.expires_at assert before + ttl <= host.expires_at <= after + ttl @@ -280,12 +279,12 @@ async def test_create_host_explicit_null_expires_at_creates_permanent_host(clien ) assert response.status_code == 201 - assert response.json()["expires_at"] is None + assert not response.json()["expires_at"] async with async_session_factory() as session: host = await session.get(Host, uuid.UUID(response.json()["id"])) - assert host is not None - assert host.expires_at is None + assert host + assert not host.expires_at async def test_create_host_rejects_oversized_idempotency_key(client): @@ -325,14 +324,11 @@ async def test_get_host_exposes_null_expires_at_by_default(client): ) assert response.status_code == 200 - assert response.json()["expires_at"] is None + assert not response.json()["expires_at"] async def test_get_host_never_returns_private_key(client): - # Regression guard: the per-VM private key is returned exactly once at - # POST time. GET must always return None — there's no column behind it, - # and any future change that accidentally persists private_key would - # break this assertion. + # The per-VM private key is returned once, at POST time; nothing persists it. host = await create_host_record(name="lb-sandbox-test", status="active") response = await client.get( @@ -341,7 +337,7 @@ async def test_get_host_never_returns_private_key(client): ) assert response.status_code == 200 - assert response.json()["private_key"] is None + assert not response.json()["private_key"] async def test_create_host_idempotency_key_returns_same_host_on_retry(client, monkeypatch): @@ -368,7 +364,6 @@ async def test_create_host_idempotency_key_returns_same_host_on_retry(client, mo assert second.status_code == 201 assert second.json()["id"] == first_id - # Only one host actually got created (kiq fires once per real create). assert mocked_provision.await_count == 1 @@ -411,28 +406,24 @@ async def test_create_host_without_idempotency_key_creates_fresh_each_time(clien async def test_idempotency_loser_in_provisioning_is_marked_for_janitor(monkeypatch): - # Race-loser that's fresh-created (in PROVISIONING) cannot be delete_host'd - # (it's in DELETE_BLOCKED_STATUSES). The cleanup path must mark it with - # expires_at=now so the janitor reaps it on the next cycle. + # delete_host refuses a provisioning host, so the loser is expired for the janitor. monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock()) async with async_session_factory() as session: service = HostService(session) - loser = await service.create_host(env={}, image=None, expires_at=None) + loser = await service.create_host(env={}, image=None, permanent=True) assert loser.status == HostStatus.PROVISIONING.value - assert loser.expires_at is None + assert not loser.expires_at await service._release_idempotency_loser(loser) async with async_session_factory() as session: reloaded = await session.get(Host, loser.id) - assert reloaded is not None - assert reloaded.expires_at is not None + assert reloaded + assert reloaded.expires_at assert reloaded.expires_at <= datetime.now(UTC) async def test_idempotency_loser_that_claimed_pool_host_is_returned_to_pool(monkeypatch): - # Race-loser that had claimed a pool host should be UN-claimed (returned - # to the pool with a fresh max-age TTL), NOT deleted. monkeypatch.setenv("POOL_SIZE", "1") monkeypatch.setenv("POOL_HOST_MAX_AGE_HOURS", "4") get_settings.cache_clear() @@ -462,42 +453,37 @@ async def test_idempotency_loser_that_claimed_pool_host_is_returned_to_pool(monk async with async_session_factory() as session: service = HostService(session) - claimed = await service._try_claim_pool_host(provider="exe", expires_at=None) - assert claimed is not None - assert claimed.claimed_at is not None + claimed = await service._try_claim_pool_host(provider="exe") + assert claimed + assert claimed.claimed_at await service._release_idempotency_loser(claimed) async with async_session_factory() as session: reloaded = await session.get(Host, pool_host.id) - assert reloaded is not None - assert reloaded.claimed_at is None - assert reloaded.expires_at is not None + assert reloaded + assert not reloaded.claimed_at + assert reloaded.expires_at assert reloaded.expires_at > datetime.now(UTC) + timedelta(hours=3) get_settings.cache_clear() async def test_record_idempotency_key_returns_false_on_duplicate(monkeypatch): - # Regression: a losing-race duplicate insert must return False without - # raising, and must leave host loaded for serialization. Recording on - # self.session expired host mid-response and tripped MissingGreenlet. monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock()) async with async_session_factory() as session: service = HostService(session) - host = await service.create_host(env={}, image=None, expires_at=None) + host = await service.create_host(env={}, image=None) - assert await service._record_idempotency_key("dup-key", host) is True - assert await service._record_idempotency_key("dup-key", host) is False + assert await service._record_idempotency_key("dup-key", host) + assert not await service._record_idempotency_key("dup-key", host) - # host stays usable — no MissingGreenlet on attribute access. - assert host.id is not None + # Recorded in self.session, host would expire before the response is built. + assert host.id assert host.status == HostStatus.PROVISIONING.value async def test_create_host_idempotency_key_expired_creates_new_host(client, monkeypatch): - # Manually insert an idempotency-key row pointing at a deleted host with - # an expired timestamp; a fresh POST with the same key should NOT reuse it. mocked_provision = AsyncMock() monkeypatch.setattr("hosts.service.HostService.provision", mocked_provision) @@ -522,7 +508,6 @@ async def test_create_host_idempotency_key_expired_creates_new_host(client, monk ) assert response.status_code == 201 - # New host was created, not the stale one. assert response.json()["id"] != str(seed_host.id) @@ -559,8 +544,8 @@ async def test_create_host_without_sizing_reflects_null(client, monkeypatch): assert response.status_code == 201 payload = response.json() - assert payload["instance_type"] is None - assert payload["disk_gb"] is None + assert not payload["instance_type"] + assert not payload["disk_gb"] async def test_create_host_rejects_sizing_the_provider_does_not_support(client): @@ -774,7 +759,7 @@ async def test_delete_host_deletes_vm_and_record_with_service_auth(client, monke assert calls == [("tailscale", "n123CNTRL"), ("exe", "lb-sandbox-test")] async with async_session_factory() as session: - assert await session.get(Host, host.id) is None + assert not await session.get(Host, host.id) async def test_delete_host_rejects_missing_or_bad_service_auth(client): @@ -859,14 +844,12 @@ async def test_delete_host_preserves_record_when_vm_teardown_fails(client, monke mocked_delete_vm.assert_awaited_once_with("lb-sandbox-test") async with async_session_factory() as session: - assert await session.get(Host, host.id) is not None + assert await session.get(Host, host.id) async def test_delete_host_skips_tailscale_release_when_device_id_missing(client, monkeypatch): - # A VM-backed host without a tailscale_device_id is treated as already - # released — delete proceeds to VM teardown and completes. This covers the - # retry path after a previous teardown that succeeded on Tailscale but - # failed before deleting the VM. + # A host without a tailscale_device_id was released already: the retry after + # a teardown that failed at the VM step. mocked_delete_device = AsyncMock() mocked_delete_vm = AsyncMock() monkeypatch.setattr("networking.tailscale.Tailscale.release_device", mocked_delete_device) @@ -886,13 +869,11 @@ async def test_delete_host_skips_tailscale_release_when_device_id_missing(client mocked_delete_vm.assert_awaited_once_with("lb-sandbox-test") async with async_session_factory() as session: - assert await session.get(Host, host.id) is None + assert not await session.get(Host, host.id) async def test_delete_host_clears_tailscale_device_id_when_vm_teardown_fails(client, monkeypatch): - # First call: release_device succeeds, delete_vm fails. The host row must - # be preserved AND have its tailscale_device_id cleared, so a retry does - # not try to delete the already-released device. + # A retry must not release the device again. mocked_delete_device = AsyncMock() mocked_delete_vm = AsyncMock(side_effect=ProviderTransportError("exe.dev API failed")) monkeypatch.setattr("networking.tailscale.Tailscale.release_device", mocked_delete_device) @@ -913,10 +894,9 @@ async def test_delete_host_clears_tailscale_device_id_when_vm_teardown_fails(cli async with async_session_factory() as session: persisted = await session.get(Host, host.id) - assert persisted is not None - assert persisted.tailscale_device_id is None + assert persisted + assert not persisted.tailscale_device_id - # Retry: delete_vm now succeeds. release_device must NOT be called again. mocked_delete_vm.side_effect = None mocked_delete_device.reset_mock() @@ -929,12 +909,10 @@ async def test_delete_host_clears_tailscale_device_id_when_vm_teardown_fails(cli mocked_delete_device.assert_not_awaited() async with async_session_factory() as session: - assert await session.get(Host, host.id) is None + assert not await session.get(Host, host.id) async def test_delete_host_succeeds_when_vm_already_absent_at_provider(client, monkeypatch): - # If exe.dev evicted the VM out from under us (or a previous delete - # partially completed), the next delete should still clean up the DB row. mocked_delete_device = AsyncMock() mocked_delete_vm = AsyncMock( side_effect=ProviderNotFoundError("vm 'lb-sandbox-test' not found") @@ -957,7 +935,7 @@ async def test_delete_host_succeeds_when_vm_already_absent_at_provider(client, m mocked_delete_vm.assert_awaited_once_with("lb-sandbox-test") async with async_session_factory() as session: - assert await session.get(Host, host.id) is None + assert not await session.get(Host, host.id) async def test_delete_host_preserves_record_when_tailscale_teardown_fails(client, monkeypatch): @@ -982,7 +960,7 @@ async def test_delete_host_preserves_record_when_tailscale_teardown_fails(client mocked_delete_vm.assert_not_awaited() async with async_session_factory() as session: - assert await session.get(Host, host.id) is not None + assert await session.get(Host, host.id) async def create_host_record( @@ -1024,9 +1002,8 @@ async def create_host_record( async def test_renew_host_with_empty_body_extends_by_default_ttl(client): - # Keepalive: an empty renew body bumps the lease to now + LEASE_DEFAULT_TTL. - # The host here is demand-provisioned (pool_member=False, claimed_at NULL), - # pinning that unclaimed non-pool hosts stay renewable by their caller. + # A demand-provisioned host has claimed_at NULL like a warm member and must + # still renew. host = await create_host_record( name="lb-renew-default", status=HostStatus.ACTIVE.value, @@ -1125,7 +1102,7 @@ async def test_renew_host_rejects_unclaimed_pool_host(client): async with async_session_factory() as session: persisted = await session.get(Host, host.id) - assert persisted is not None + assert persisted assert persisted.expires_at == host.expires_at @@ -1199,8 +1176,6 @@ async def test_renew_host_rejects_missing_or_bad_service_auth(client): async def test_delete_host_maps_database_failure_to_503(client, monkeypatch): - # A SQLAlchemyError during teardown must surface as a controlled 503, like - # create_host's DB-failure path, not an unhandled 500. monkeypatch.setattr( "hosts.service.HostService.delete_host", AsyncMock(side_effect=SQLAlchemyError("db down")), diff --git a/src/hosts/tests/test_max_lifetime.py b/src/hosts/tests/test_max_lifetime.py index 0f62645..48949c8 100644 --- a/src/hosts/tests/test_max_lifetime.py +++ b/src/hosts/tests/test_max_lifetime.py @@ -79,7 +79,7 @@ async def test_pool_leases_end_when_the_provider_stops_the_vm(limited_provider): late = stop + timedelta(minutes=1) assert not await service._try_claim_pool_host(provider="stub", expires_at=late) - claimed = await service._try_claim_pool_host(provider="stub", expires_at=...) + claimed = await service._try_claim_pool_host(provider="stub") assert claimed assert claimed.expires_at == stop diff --git a/src/hosts/tests/test_models.py b/src/hosts/tests/test_models.py index fb24f25..0f868b3 100644 --- a/src/hosts/tests/test_models.py +++ b/src/hosts/tests/test_models.py @@ -60,15 +60,11 @@ async def _patch_provision_happy_path(monkeypatch, *, ssh_port: int = 22) -> dic def test_build_name_is_distinct_for_uuids_sharing_a_timestamp_prefix(): - # Regression: UUIDv7's first 48 bits are the millisecond timestamp, - # so two concurrent creates produce host_ids whose leading hex is - # identical. build_name() must derive from the random tail, not the - # timestamp prefix, or it collides on the unique name index. + # UUIDv7 leads with a millisecond timestamp, so two concurrent creates share + # their leading hex; the name must come from the random tail. a = uuid.UUID("019e8307-79af-7000-8000-000000000001") b = uuid.UUID("019e8307-79af-7000-8000-000000000002") - # Sanity: the offending prefix really is identical. assert a.hex[:12] == b.hex[:12] - # And the names build_name produces must not be. assert Host.build_name(a) != Host.build_name(b) @@ -93,37 +89,36 @@ async def test_provision_walks_host_to_active(monkeypatch): async with async_session_factory() as session: refreshed = await session.get(Host, host.id) - assert refreshed is not None + assert refreshed assert refreshed.name == "exe-runtime-1" assert refreshed.internal_ssh_host == "exe-runtime-1.example.ts.net" assert refreshed.external_ssh_host == "" assert refreshed.status == "active" assert refreshed.tailscale_device_id == "n123CNTRL" assert refreshed.known_hosts == "exe-runtime-1.example.ts.net ssh-ed25519 AAAATEST\n" - assert refreshed.activated_at is not None + assert refreshed.activated_at mocks["issue_join_credentials"].assert_awaited_once_with(host_name="lb-sandbox-test") - assert mocks["create_vm"].await_args is not None + assert mocks["create_vm"].await_args create_vm_kwargs = mocks["create_vm"].await_args.kwargs assert create_vm_kwargs["name"] == "lb-sandbox-test" assert create_vm_kwargs["image"] == "ghcr.io/drukbox/custom-sandbox:provision-test" # No per-request sizing on the row → the provider falls back to its # configured default size. - assert create_vm_kwargs["instance_type"] is None - assert create_vm_kwargs["disk_gb"] is None + assert not create_vm_kwargs["instance_type"] + assert not create_vm_kwargs["disk_gb"] assert create_vm_kwargs["env"]["TAILSCALE_AUTHKEY"] == "tskey-secret" assert create_vm_kwargs["env"]["TAILSCALE_ADVERTISE_TAGS"] == "tag:sandbox" assert create_vm_kwargs["env"]["SANDBOX_GATEWAY_URL"] == "wss://gateway.example.ts.net/daemon" # No drukbox-side callbacks injected — the box never phones home. assert "REMOTE_HOST_ANNOUNCE_URL" not in create_vm_kwargs["env"] assert "REMOTE_HOST_ANNOUNCE_TOKEN" not in create_vm_kwargs["env"] - # Bootstrap script ships to the VM only knows about Tailscale. assert "TAILSCALE_AUTHKEY" in create_vm_kwargs["setup_script"] assert "REMOTE_HOST_ANNOUNCE" not in create_vm_kwargs["setup_script"] - # wait_for_device is called against the post-create VM name (mirrors how - # ExeProvider can rename hosts during creation). + # A provider may rename the VM during creation; discovery uses the name it + # returned. mocks["wait_for_device"].assert_awaited_once() - assert mocks["wait_for_device"].await_args is not None + assert mocks["wait_for_device"].await_args wait_kwargs = mocks["wait_for_device"].await_args.kwargs assert wait_kwargs["host_name"] == "exe-runtime-1" @@ -143,17 +138,14 @@ async def test_provision_forwards_sizing_from_the_row_to_create_vm(monkeypatch): session.add(host) await HostService(session).provision(host) - assert mocks["create_vm"].await_args is not None + assert mocks["create_vm"].await_args create_vm_kwargs = mocks["create_vm"].await_args.kwargs assert create_vm_kwargs["instance_type"] == "t3.xlarge" assert create_vm_kwargs["disk_gb"] == 250 async def test_provision_threads_ssh_username_from_vm_result_onto_host(monkeypatch): - # The provider knows which in-VM user the image runs as; provision() - # must propagate it so the HostOut response (and subsequent GETs) - # carry the right value without callers having to maintain their own - # per-provider mapping. + # The provider knows which user the image runs as. mocks = await _patch_provision_happy_path(monkeypatch) mocks["create_vm"].return_value = VMCreateResult( provider_id="exe-runtime-1", @@ -163,15 +155,14 @@ async def test_provision_threads_ssh_username_from_vm_result_onto_host(monkeypat ) async with async_session_factory() as session: - host = await HostService(session).create_host(env={}, image=None, expires_at=None) + host = await HostService(session).create_host(env={}, image=None) assert host.ssh_username == "exedev" async def test_provision_threads_private_key_from_vm_result_onto_host(monkeypatch): - # Providers that mint a per-VM keypair (AWS) return private key material - # in VMCreateResult; provision() must propagate it onto the host instance - # so the create-time HostOut response carries it. + # A provider that issues a per-VM keypair returns the private key; the POST + # response carries it. mocks = await _patch_provision_happy_path(monkeypatch) mocks["create_vm"].return_value = VMCreateResult( provider_id="exe-runtime-1", @@ -182,22 +173,21 @@ async def test_provision_threads_private_key_from_vm_result_onto_host(monkeypatc ) async with async_session_factory() as session: - host = await HostService(session).create_host(env={}, image=None, expires_at=None) + host = await HostService(session).create_host(env={}, image=None) - assert host.private_key is not None + assert host.private_key assert host.private_key.startswith("-----BEGIN OPENSSH PRIVATE KEY-----") async def test_create_host_clears_safety_ttl_when_caller_omits_one(monkeypatch): - # The row is born with a safety TTL so the janitor reaps it if the - # client disconnects mid-provision. On success, that safety TTL is - # replaced with the caller's intent — here, NULL. + # The safety TTL the row is born with gives way to the caller's lease, here + # none. await _patch_provision_happy_path(monkeypatch) async with async_session_factory() as session: - created = await HostService(session).create_host(env={}, image=None, expires_at=None) + created = await HostService(session).create_host(env={}, image=None, permanent=True) - assert created.expires_at is None + assert not created.expires_at async def test_create_host_lands_on_caller_ttl_on_success(monkeypatch): @@ -207,14 +197,13 @@ async def test_create_host_lands_on_caller_ttl_on_success(monkeypatch): async with async_session_factory() as session: created = await HostService(session).create_host(env={}, image=None, expires_at=caller_ttl) - assert created.expires_at is not None + assert created.expires_at assert abs((created.expires_at - caller_ttl).total_seconds()) < 1 async def test_create_host_expires_errored_host_for_janitor(monkeypatch): - # Strand protection on the unhappy path: if provisioning errors out, - # expire the row immediately so the janitor reaps it (and its VM) on the - # next cycle, without operators having to chase ERROR hosts manually. + # A failed provision expires the row at once so the janitor reaps it and its + # VM. mocks = await _patch_provision_happy_path(monkeypatch) mocks["wait_for_device"].side_effect = DeviceDiscoveryTimeoutError( "Tailscale never surfaced device" @@ -222,18 +211,13 @@ async def test_create_host_expires_errored_host_for_janitor(monkeypatch): async with async_session_factory() as session: service = HostService(session) - try: - await service.create_host(env={}, image=None, expires_at=None) - except ProvisioningFailedError: - pass - else: - raise AssertionError("expected ProvisioningFailedError") - - # Row should still exist, expired so the janitor reaps it next cycle. + with pytest.raises(ProvisioningFailedError): + await service.create_host(env={}, image=None) + async with async_session_factory() as session: result = await session.execute(select(Host).where(Host.status == "error")) errored = result.scalar_one() - assert errored.expires_at is not None + assert errored.expires_at assert errored.expires_at <= utc_now() @@ -252,7 +236,7 @@ async def test_provision_generated_authkey_overrides_caller_env(monkeypatch): session.add(host) await HostService(session).provision(host) - assert mocks["create_vm"].await_args is not None + assert mocks["create_vm"].await_args assert mocks["create_vm"].await_args.kwargs["env"]["TAILSCALE_AUTHKEY"] == "tskey-secret" @@ -267,7 +251,7 @@ async def test_provision_marks_error_when_provider_rejects_credentials(monkeypat async with async_session_factory() as session: refreshed = await session.get(Host, host.id) - assert refreshed is not None + assert refreshed assert refreshed.status == "error" assert refreshed.last_error == "ProviderAuthError: exe.dev authentication failed" @@ -285,10 +269,9 @@ async def test_provision_marks_error_when_device_discovery_times_out(monkeypatch async with async_session_factory() as session: refreshed = await session.get(Host, host.id) - assert refreshed is not None + assert refreshed assert refreshed.status == "error" assert "Tailscale did not surface device" in refreshed.last_error - # ssh-keyscan must not have been attempted if Tailscale never saw the device. mocks["scan_known_hosts"].assert_not_awaited() @@ -303,7 +286,7 @@ async def test_provision_marks_error_when_keyscan_fails(monkeypatch): async with async_session_factory() as session: refreshed = await session.get(Host, host.id) - assert refreshed is not None + assert refreshed assert refreshed.status == "error" assert "keyscan failed" in refreshed.last_error # Device id was already stored when Tailscale surfaced it; teardown @@ -312,9 +295,7 @@ async def test_provision_marks_error_when_keyscan_fails(monkeypatch): async def test_scan_known_hosts_retries_until_banner_is_returned(monkeypatch): - # Mirrors the production race: tailscaled-SSH is not ready immediately - # when the device shows up in Tailscale's API. First keyscan attempt - # gets nothing back; a subsequent attempt succeeds. + # tailscaled-SSH is not ready the moment the device shows up in the API. host = Host( name="lb-sandbox-test", image=ExeSettings().default_image, # pyright: ignore[reportCallIssue] @@ -435,7 +416,6 @@ async def test_scan_known_hosts_retries_when_only_one_target_surfaces(monkeypatc returncode=0, communicate=AsyncMock(return_value=(b"sb-test.exe.xyz ssh-rsa AAAEXE\n", b"")), ) - # Attempt 1: internal answers, external sshd still warming. Attempt 2: both. create_subprocess_exec = AsyncMock( side_effect=[internal_key, external_warming, internal_key, external_key] ) @@ -481,7 +461,7 @@ async def test_scan_known_hosts_uses_async_subprocess(monkeypatch): assert known_hosts == b"lb-sandbox-test.example.ts.net ssh-ed25519 AAAATEST\n" create_subprocess_exec.assert_awaited_once() - assert create_subprocess_exec.await_args is not None + assert create_subprocess_exec.await_args assert create_subprocess_exec.await_args.args == ( "ssh-keyscan", "-p", @@ -491,7 +471,6 @@ async def test_scan_known_hosts_uses_async_subprocess(monkeypatch): async def test_scan_known_hosts_honors_external_ssh_port(monkeypatch): - """Scans the external address on its published port, not a hardcoded 22.""" host = Host( name="sb-test", image=ExeSettings().default_image, # pyright: ignore[reportCallIssue] @@ -517,7 +496,7 @@ async def test_scan_known_hosts_honors_external_ssh_port(monkeypatch): assert b"127.0.0.1" in known_hosts create_subprocess_exec.assert_awaited_once() - assert create_subprocess_exec.await_args is not None + assert create_subprocess_exec.await_args assert create_subprocess_exec.await_args.args == ("ssh-keyscan", "-p", "49160", "127.0.0.1") @@ -579,5 +558,5 @@ async def test_expires_at_with_non_utc_offset_round_trips_as_utc(): async with async_session_factory() as session: reloaded = await session.get(Host, host.id) - assert reloaded is not None + assert reloaded assert reloaded.expires_at == datetime(2099, 1, 1, 8, 0, tzinfo=UTC) diff --git a/src/hosts/tests/test_pool.py b/src/hosts/tests/test_pool.py index ac8d19a..0b12a36 100644 --- a/src/hosts/tests/test_pool.py +++ b/src/hosts/tests/test_pool.py @@ -21,9 +21,7 @@ def pooled_settings(monkeypatch): monkeypatch.setenv("POOL_SIZE", "2") monkeypatch.delenv("POOL_SIZES", raising=False) monkeypatch.setenv("POOL_HOST_MAX_AGE_HOURS", "4") - # Raise the per-tick cap so tests can observe full deficit refills in a - # single maintain_pool() call. Production defaults to a small cap so that - # overlapping ticks bound the over-provision blast radius. + # A high per-tick cap lets one maintain_pool() call refill a whole deficit. monkeypatch.setenv("POOL_MAX_CREATES_PER_TICK", "100") get_settings.cache_clear() yield get_settings() @@ -99,15 +97,13 @@ async def test_create_host_claims_pool_member_when_eligible(pooled_settings, mon result = await service.get_or_create_host(env={}, image=None) assert result.id == pool_host.id - assert result.claimed_at is not None + assert result.claimed_at mocked_provision.assert_not_awaited() async def test_pool_never_claims_or_sheds_demand_hosts(pooled_settings, monkeypatch): - # A demand-provisioned host keeps claimed_at NULL just like a warm pool - # host, but pool_member=False must keep pool claim/count/shed away from it — - # otherwise caller A's env-bearing sandbox could be handed to caller B, or - # deleted as pool excess. + # A demand host has claimed_at NULL like a warm member; only pool_member keeps + # the claim, count and shed away from it. demand = await _seed_pool_host( name="lb-demand-1", pool_member=False, env={"SECRET": "caller-a"} ) @@ -115,17 +111,15 @@ async def test_pool_never_claims_or_sheds_demand_hosts(pooled_settings, monkeypa monkeypatch.setattr("networking.tailscale.Tailscale.release_device", AsyncMock()) monkeypatch.setattr("providers.exe.provider.ExeProvider.delete_vm", AsyncMock()) - # A later default request must fall back to fresh, not claim the demand host. async with async_session_factory() as session: service = HostService(session, settings=pooled_settings) result = await service.get_or_create_host(env={}, image=None) assert result.id != demand.id - assert result.claimed_at is None + assert not result.claimed_at - # Pool maintenance must not count or shed it; the demand host survives. await maintain_pool() async with async_session_factory() as session: - assert await session.get(Host, demand.id) is not None + assert await session.get(Host, demand.id) async def test_create_host_falls_back_to_fresh_when_pool_empty(pooled_settings, monkeypatch): @@ -136,7 +130,7 @@ async def test_create_host_falls_back_to_fresh_when_pool_empty(pooled_settings, service = HostService(session, settings=pooled_settings) result = await service.get_or_create_host(env={}, image=None) - assert result.claimed_at is None + assert not result.claimed_at assert result.status == HostStatus.PROVISIONING.value mocked_provision.assert_awaited_once() @@ -150,7 +144,7 @@ async def test_create_host_skips_pool_when_env_present(pooled_settings, monkeypa service = HostService(session, settings=pooled_settings) result = await service.get_or_create_host(env={"SANDBOX_X": "value"}, image=None) - assert result.claimed_at is None + assert not result.claimed_at mocked_provision.assert_awaited_once() @@ -168,7 +162,7 @@ async def test_create_host_skips_pool_when_provider_has_no_target( service = HostService(session, settings=pooled_settings) result = await service.get_or_create_host(env={}, image=None, provider="stub") - assert result.claimed_at is None + assert not result.claimed_at assert result.provider == "stub" mocked_provision.assert_awaited_once() @@ -186,7 +180,7 @@ async def test_create_host_claims_pool_member_of_the_pinned_provider( result = await service.get_or_create_host(env={}, image=None, provider="stub") assert result.id == stub_host.id - assert result.claimed_at is not None + assert result.claimed_at mocked_provision.assert_not_awaited() @@ -205,7 +199,7 @@ async def test_create_host_never_claims_another_providers_pool_member( assert result.id != exe_host.id assert result.provider == "stub" - assert result.claimed_at is None + assert not result.claimed_at mocked_provision.assert_awaited_once() @@ -249,7 +243,7 @@ async def test_create_host_skips_pool_when_sized(multi_pool_settings, monkeypatc ) assert result.id != pool_host.id - assert result.claimed_at is None + assert not result.claimed_at assert result.instance_type == instance_type assert result.disk_gb == disk_gb @@ -265,7 +259,7 @@ async def test_create_host_skips_pool_when_image_override(pooled_settings, monke service = HostService(session, settings=pooled_settings) result = await service.get_or_create_host(env={}, image="custom-image:tag") - assert result.claimed_at is None + assert not result.claimed_at assert result.image == "custom-image:tag" mocked_provision.assert_awaited_once() @@ -282,7 +276,7 @@ async def test_pool_claim_overrides_pool_max_age_with_caller_expires_at( service = HostService(session, settings=pooled_settings) claimed = await service.get_or_create_host(env={}, image=None, expires_at=caller_expires_at) - assert claimed.expires_at is not None + assert claimed.expires_at assert abs((claimed.expires_at - caller_expires_at).total_seconds()) < 1 @@ -298,24 +292,22 @@ async def test_pool_claim_stamps_default_lease_when_caller_omits_expires_at( async with async_session_factory() as session: service = HostService(session, settings=pooled_settings) - claimed = await service.get_or_create_host(env={}, image=None) # no expires_at + claimed = await service.get_or_create_host(env={}, image=None) lease_end = utc_now() + timedelta(seconds=pooled_settings.lease_default_ttl) - assert claimed.expires_at is not None + assert claimed.expires_at assert abs((claimed.expires_at - lease_end).total_seconds()) < 5 async def test_pool_claim_honors_explicit_permanent_lease(pooled_settings, monkeypatch): - # An explicit expires_at=None is the caller's deliberate opt-in to a - # permanent host; the claim must clear the warm-pool max-age TTL. await _seed_pool_host(name="lb-pool-1", expires_at=datetime.now(UTC) + timedelta(hours=4)) monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock()) async with async_session_factory() as session: service = HostService(session, settings=pooled_settings) - claimed = await service.get_or_create_host(env={}, image=None, expires_at=None) + claimed = await service.get_or_create_host(env={}, image=None, permanent=True) - assert claimed.expires_at is None + assert not claimed.expires_at async def test_concurrent_claims_dont_grab_the_same_pool_host(pooled_settings, monkeypatch): @@ -330,8 +322,8 @@ async def attempt_claim(): first, second = await asyncio.gather(attempt_claim(), attempt_claim()) - claim_ids = {h.id for h in (first, second) if h.claimed_at is not None} - fresh_ids = {h.id for h in (first, second) if h.claimed_at is None} + claim_ids = {host.id for host in (first, second) if host.claimed_at} + fresh_ids = {host.id for host in (first, second) if not host.claimed_at} assert claim_ids == {pool_host.id} assert len(fresh_ids) == 1 mocked_provision.assert_awaited_once() @@ -358,7 +350,6 @@ async def test_maintain_pool_tops_up_only_the_deficit(pooled_settings, monkeypat async def test_maintain_pool_sheds_excess_when_pool_size_shrinks(pooled_settings, monkeypatch): - # POOL_SIZE=2; seed 5 unclaimed hosts. Maintainer should shed the 3 oldest. monkeypatch.setattr("networking.tailscale.Tailscale.release_device", AsyncMock()) monkeypatch.setattr("providers.exe.provider.ExeProvider.delete_vm", AsyncMock()) for i in range(5): @@ -484,8 +475,8 @@ async def create_host_with_exe_down(self, **kwargs): async def test_maintain_pool_excludes_expired_hosts_from_count(pooled_settings, monkeypatch): - # An unclaimed host past its expires_at is queued for janitor reaping and - # must not count toward the pool target. Maintainer should top up to N. + # An unclaimed host past its expires_at is left to the janitor and does not + # count toward the target. monkeypatch.setattr("hosts.service.HostService.provision", AsyncMock()) await _seed_pool_host( name="lb-pool-expiring", @@ -494,7 +485,6 @@ async def test_maintain_pool_excludes_expired_hosts_from_count(pooled_settings, summary = await maintain_pool() - # The expiring host doesn't count, so the maintainer creates pool_size more. assert summary.created == pooled_settings.pool_size @@ -520,7 +510,6 @@ async def test_maintain_pool_breaks_after_consecutive_create_failures(pooled_set provision = AsyncMock(side_effect=RuntimeError("provider down")) monkeypatch.setattr("hosts.service.HostService.provision", provision) - # Force pool_size large enough that the cap matters. monkeypatch.setenv("POOL_SIZE", "10") get_settings.cache_clear() try: @@ -534,7 +523,6 @@ async def test_maintain_pool_breaks_after_consecutive_create_failures(pooled_set async def test_create_host_ignores_pool_when_pool_size_is_zero(monkeypatch): - # Default settings: pool_size = 0. get_settings.cache_clear() await _seed_pool_host(name="ch-orphan-pool-host") mocked_provision = AsyncMock() @@ -544,7 +532,7 @@ async def test_create_host_ignores_pool_when_pool_size_is_zero(monkeypatch): service = HostService(session) result = await service.get_or_create_host(env={}, image=None) - assert result.claimed_at is None + assert not result.claimed_at mocked_provision.assert_awaited_once() @@ -562,4 +550,4 @@ async def test_pool_shed_skips_a_host_claimed_in_the_race(pooled_settings, monke assert not deleted delete_vm.assert_not_awaited() async with async_session_factory() as session: - assert await session.get(Host, claimed.id) is not None + assert await session.get(Host, claimed.id)