diff --git a/src/agent_env/env/env.py b/src/agent_env/env/env.py index 02b97788..f19312bd 100644 --- a/src/agent_env/env/env.py +++ b/src/agent_env/env/env.py @@ -74,6 +74,7 @@ class DeployedEnv: # From the card whenever the record carries it (__post_init__); stored because readers of the JSON use them. mcp_url: Optional[str] = None mcp_server_name: Optional[str] = None + preloaded_universe: Optional[dict] = None # Set by the registry and the deploy step, never by a provider. metadata: Optional[dict] = None instance_id: Optional[str] = None @@ -134,6 +135,7 @@ def _fields_from(cls, data: dict) -> dict: instance_id=data.get("instance_id"), created_at_utc=data.get("created_at_utc"), expires_at_utc=data.get("expires_at_utc"), + preloaded_universe=data.get("preloaded_universe"), ) @@ -188,6 +190,7 @@ def _fields_from(cls, data: dict) -> dict: @dataclass class LoadEnvironmentUniverseArtifactResult: metadata_filepaths: dict[str, str] = field(default_factory=dict) + delivered_preloaded: bool = False # Which path the load actually took. A snapshot restore swaps a pre-ingested servicedb # image; a miss re-ingests every service over HTTP under a 600s per-service cap. The # two are comparable in wall-clock on a well-sized VM (measured ~4min vs ~3.5min on a diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 5fdaf27a..a077dbb2 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -157,6 +157,7 @@ async def load_environment_artifact(self, environment_artifact: EnvironmentArtif # A built-in's deploy, or a caller wiring one up by hand, leaves a gateway; a plugin's leaves only its record. if self._gateway_url is None and self._deployed is None: raise RuntimeError("Environment not deployed - call deploy() first") + self._void_preloaded_universe() # MCP servers are checked first; if a website shares the same environment_name, the MCP server wins. for env in self.mcp_server_envs: if env.environment_name == environment_artifact.environment_name: @@ -172,6 +173,7 @@ async def load_environment_universe_artifact( self, environment_universe_artifact: EnvironmentUniverseArtifact, snapshot_after_load: bool | None = None, + skip_if_preloaded: bool = False, ) -> LoadEnvironmentUniverseArtifactResult: """Seed this env from a universe artifact, preferring a pre-baked snapshot. @@ -192,6 +194,10 @@ async def load_environment_universe_artifact( and pushes a multi-GB servicedb image (~5min, ~2.3GB stored) — worth it to make a universe cheap to re-serve, not worth it on a one-off deploy. The bake is best-effort and never fails a load that already succeeded. + + ``skip_if_preloaded`` skips the data load for a deployment its provider delivered + with this universe already loaded and untouched since (see ``_delivered_with``). It + is off by default, so a caller's load always loads unless the caller asks otherwise. """ import uuid from agent_env.env.env import LoadEnvironmentUniverseArtifactResult @@ -212,17 +218,26 @@ async def load_environment_universe_artifact( # Keyed on the service images rather than self.id, so re-registering the same # servers under a new env id still hits (and a rebuilt server correctly misses). env_fingerprint = compute_env_fingerprint(self) if can_restore else None + preloaded = skip_if_preloaded and await self._delivered_with(environment_universe_artifact) + if not preloaded: + self._void_preloaded_universe() snapshot = get_env_snapshot_store().get_clean( self.id, environment_universe_artifact.id, environment_universe_artifact.version, env_fingerprint=env_fingerprint, - ) if can_restore else None + ) if can_restore and not preloaded else None restored_from_snapshot = False snapshot_baked: bool | None = None snapshot_bake_error: str | None = None - if snapshot: + if preloaded: + logger.info( + "env=%s was delivered with universe %s v%s already loaded and untouched; " + "skipping the data load", + self.id, environment_universe_artifact.id, environment_universe_artifact.version, + ) + elif snapshot: logger.info(f"Found clean snapshot for env={self.id} universe={environment_universe_artifact.id}: artifact={snapshot.db_image_artifact_id} v{snapshot.db_image_artifact_version} instance={snapshot.instance_id}") await self._load_from_snapshot(snapshot) restored_from_snapshot = True @@ -334,7 +349,7 @@ async def _load_one(idx: int, artifact): # image swap. Deliberately placed here: after the universe stamp, which # EnvSnapshot.create reads back off the instance record, and before the metadata # download, so the captured PGDATA reflects the data load and nothing else. - if not restored_from_snapshot and can_restore: + if not restored_from_snapshot and not preloaded and can_restore: if snapshot_after_load is None: snapshot_after_load = _snapshot_after_load_default() if snapshot_after_load: @@ -361,12 +376,60 @@ async def _load_one(idx: int, artifact): logger.info(f"Successfully downloaded metadata '{key}'") return LoadEnvironmentUniverseArtifactResult( metadata_filepaths=metadata_filepaths, + delivered_preloaded=preloaded, restored_from_snapshot=restored_from_snapshot, snapshot_db_image_artifact_id=snapshot.db_image_artifact_id if snapshot else None, snapshot_baked=snapshot_baked, snapshot_bake_error=snapshot_bake_error, ) + async def _delivered_with(self, environment_universe_artifact) -> bool: + """Whether this deployment was handed over with this universe already loaded, and + nothing has touched its data since. + + Only the provider that made the deployment can say so, by setting + ``preloaded_universe`` on the record it returns: a warm pool, say, that loaded the + universe before anyone asked for it. This is what lets a complete load be skipped + where resume never skips one. Resume can't tell a caller's explicit load from a + retry, but here the deployment itself guarantees the data is the universe asked + for. The changelog gate still applies, so data touched since breaks the guarantee + and the universe loads as usual. + """ + preloaded = getattr(self._deployed, "preloaded_universe", None) or {} + if (preloaded.get("id"), preloaded.get("version")) != ( + environment_universe_artifact.id, environment_universe_artifact.version, + ): + return False + if self._sandbox is None or self._sandbox.mode == SANDBOX_MODE_CONTAINER: + return False + try: + from agent_env.env.snapshot_store import _check_changelog_empty + + return await _check_changelog_empty(self._sandbox) + except Exception as e: # noqa: BLE001 + logger.warning("Could not confirm the preloaded universe is untouched (%s); loading it", e) + return False + + def _void_preloaded_universe(self) -> None: + """Drop the deployment's preloaded guarantee before a load replaces its data. + + A load reinstalls each service's changelog triggers, which empties the changelog. + Left standing, a guarantee for universe A would survive loading B and pass the + changelog gate on the next request for A, skipping a load the data needs. Cleared on + the stored record too, so a reattach can't revive it. + """ + if not getattr(self._deployed, "preloaded_universe", None): + return + self._deployed.preloaded_universe = None + if not self._instance_id: + return + try: + from agent_env.env.store import get_env_instance_store + + get_env_instance_store().clear_preloaded_universe(self._instance_id) + except Exception as e: # noqa: BLE001 + logger.warning("Could not void the preloaded universe on instance %s: %s", self._instance_id, e) + async def _resumable_environments(self, environment_universe_artifact, total: int) -> set[str]: """Services safe to skip because a previous attempt already loaded them. diff --git a/src/agent_env/env/store.py b/src/agent_env/env/store.py index 24e83ffe..cf12f392 100644 --- a/src/agent_env/env/store.py +++ b/src/agent_env/env/store.py @@ -300,6 +300,14 @@ def clear_loaded_environments(self, instance_id: str) -> None: UpdateSpec(set={"loaded_environments": {}}), ) + def clear_preloaded_universe(self, instance_id: str) -> None: + """Void the universe an instance was delivered with, once a load replaces its data.""" + self._doc_store.update( + ENV_INSTANCES_COLLECTION, + Filter.of(instance_id=instance_id), + UpdateSpec(set={"preloaded_universe": None}), + ) + def merge_metadata(self, instance_id: str, values: dict) -> None: """Merge key/values into an instance record's ``metadata``. diff --git a/src/agent_env/providers/env_providers/env_provider.py b/src/agent_env/providers/env_providers/env_provider.py index 16dd2982..76de8344 100644 --- a/src/agent_env/providers/env_providers/env_provider.py +++ b/src/agent_env/providers/env_providers/env_provider.py @@ -33,6 +33,7 @@ class EnvironmentProvider(ABC): """Deploys an env and hands back its record; tears down what it created.""" type: ClassVar[str] # the name an env declares it by, and that its records carry as env_provider_type + delivers_universe: ClassVar[bool] = False @abstractmethod async def deploy(self, env: Env, sandbox_provider: SandboxProvider) -> DeployedEnv: diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index eb53483a..694ca38a 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -120,6 +120,29 @@ def preflight(self) -> list[str]: return [f"deploy_env '{self.id}': env '{self.env_id}' can't be loaded: {e}"] return [f"deploy_env '{self.id}': {refusal}"] if refusal else [] + async def _load_delivered_universe(self, env, context: TaskStepContext) -> None: + """Load this step's universe when the env's provider promises a loaded env. + + Such a provider (``delivers_universe``) makes deploy_env with an artifact a contract: + the env has that universe whether the provider handed over one already holding it or + built one cold. Its task carries no load_artifact step, so the load happens here, + skipped when the env arrived with this universe untouched, and is recorded the way + load_artifact records it. Every other provider leaves loading to load_artifact. + """ + from agent_env.artifact import Artifact, EnvironmentUniverseArtifact + from agent_env.env.envs.multi_env import MultiEnv + from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep + + if self.artifact_id is None or not isinstance(env, MultiEnv): + return + if not getattr(getattr(env, "_env_provider", None), "delivers_universe", False): + return + artifact = Artifact.get(self.artifact_id, self.artifact_version) + if not isinstance(artifact, EnvironmentUniverseArtifact): + return + load_result = await env.load_environment_universe_artifact(artifact, skip_if_preloaded=True) + LoadArtifactTaskStep._record_universe_load(context, artifact, load_result) + async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.env.env import Env @@ -180,6 +203,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if deployed_env.instance_id: update_env_instance_metadata(deployed_env.instance_id, annotations) context.deployed_envs.append(deployed_env) + await self._load_delivered_universe(env, context) # Some envs (e.g. harbor coding tasks) build + run the agent as part of # deploy() and have no separate DeployAgent step. They advertise the diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index a0c64e93..a2c4f65a 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -267,6 +267,8 @@ def _record_universe_load(context: TaskStepContext, artifact, load_result) -> No value = getattr(load_result, key, None) if value is not None: entry[key] = value + if getattr(load_result, "delivered_preloaded", False): + entry["delivered_preloaded"] = True context.metadata.setdefault("loaded_environment_universes", []).append(entry) @staticmethod diff --git a/tst/unit/env/envs/test_multi_env_preloaded_universe.py b/tst/unit/env/envs/test_multi_env_preloaded_universe.py new file mode 100644 index 00000000..53b50c91 --- /dev/null +++ b/tst/unit/env/envs/test_multi_env_preloaded_universe.py @@ -0,0 +1,251 @@ +"""A deployment handed over with its universe already loaded skips the data load. + +A provider can deliver an env whose universe was loaded before anyone asked for it -- a +warm pool, say. Re-ingesting it would cost the whole load the provider exists to save. +Resume never skips a complete load, because it can't tell a caller's explicit load from a +retry; here the deployment's own record guarantees the data is the universe asked for. + +The safety property is the same changelog gate resume uses: a guarantee about data that +has since been touched is void, and the universe loads as usual. +""" + +from __future__ import annotations + +import dataclasses +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agent_env.env.env import DeployedEnv, DeployedGatewayEnv +from agent_env.env.envs.multi_env import MultiEnv +from agent_env.providers.env_providers import EnvironmentGatewayProvider +from agent_env.providers.env_state import LocalPostgresStateProvider +from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER + +SERVICES = ["a", "b", "gmail"] + + +def _universe(version: int = 39, metadata: dict | None = None): + ua = MagicMock() + ua.id, ua.version = "uni", version + ua.get_metadata.return_value = metadata or {} + arts = [] + for name in SERVICES: + a = MagicMock() + a.environment_name = name + arts.append(a) + ua.get_environment_artifacts.return_value = arts + return ua + + +def _env(*, preloaded: dict | None, mode: str = "vm"): + env = MultiEnv(id="env-1", version=1, mcp_server_envs=[]) + env._instance_id = "inst-1" + env._deployed = DeployedGatewayEnv( + env_id="env-1", env_version=1, instance_id="inst-1", sandbox_id="sb-1", + gateway_url="https://gw.example", preloaded_universe=preloaded, + ) + gp = EnvironmentGatewayProvider() + gp._state_provider = LocalPostgresStateProvider() + env._env_provider = gp + sandbox = MagicMock() + sandbox.mode = mode + sandbox.exec_script = AsyncMock(return_value="8\n") + sandbox.load_object_file = AsyncMock() + env._sandbox = sandbox + env._load_from_snapshot = AsyncMock() + env._bake_snapshot_after_load = AsyncMock(return_value=(True, None)) + return env + + +async def _load(env, universe, *, changelog_empty=True, changelog_error=None, skip_if_preloaded=True): + """Drive a load with stubbed stores and changelog; returns what it touched.""" + loaded: list[str] = [] + env.load_environment_artifact = AsyncMock(side_effect=lambda a: loaded.append(a.environment_name)) + instances = MagicMock() + instances.get_loaded_environments.return_value = [] + snapshots = MagicMock() + snapshots.get_clean.return_value = None + stamp = MagicMock() + changelog = AsyncMock(side_effect=changelog_error) if changelog_error else AsyncMock(return_value=changelog_empty) + with patch("agent_env.env.store.get_env_instance_store", return_value=instances), \ + patch("agent_env.env.snapshot_store.get_env_snapshot_store", return_value=snapshots), \ + patch("agent_env.env.snapshot_store._check_changelog_empty", changelog), \ + patch("agent_env.env.store.update_env_instance_environment_universe", stamp): + result = await env.load_environment_universe_artifact( + universe, snapshot_after_load=True, skip_if_preloaded=skip_if_preloaded, + ) + return MagicMock(loaded=loaded, snapshots=snapshots, stamp=stamp, result=result, instances=instances) + + +@pytest.mark.asyncio +async def test_a_deployment_delivered_with_the_universe_loads_no_service(): + """The saving: the whole re-ingest the provider loaded ahead of time.""" + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe()) + assert run.loaded == [] + + +@pytest.mark.asyncio +async def test_the_universe_is_still_recorded_on_the_instance(): + """Readers of the instance record ask what it holds, however it came to hold it.""" + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe()) + run.stamp.assert_called_once_with("inst-1", "uni", 39) + + +@pytest.mark.asyncio +async def test_metadata_files_are_still_staged(): + """Later steps read the paths this returns; they live under a fresh directory per load, + so the provider's earlier load can't have produced them.""" + file_artifact = MagicMock(filename="universe.yaml", object_url="s3://b/universe.yaml") + env = _env(preloaded={"id": "uni", "version": 39}) + run = await _load(env, _universe(metadata={"universe": file_artifact})) + env._sandbox.load_object_file.assert_awaited_once() + assert run.result.metadata_filepaths["universe"].endswith("/universe.yaml") + + +@pytest.mark.asyncio +async def test_a_preloaded_env_neither_restores_nor_bakes_a_snapshot(): + env = _env(preloaded={"id": "uni", "version": 39}) + run = await _load(env, _universe()) + run.snapshots.get_clean.assert_not_called() + env._load_from_snapshot.assert_not_awaited() + env._bake_snapshot_after_load.assert_not_awaited() + assert run.result.restored_from_snapshot is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("preloaded", [ + None, + {"id": "other", "version": 39}, + {"id": "uni", "version": 38}, + {}, +]) +async def test_a_deployment_not_delivered_with_this_universe_loads_it(preloaded): + """Another universe, or another version of this one, is not this universe.""" + run = await _load(_env(preloaded=preloaded), _universe()) + assert sorted(run.loaded) == sorted(SERVICES) + + +@pytest.mark.asyncio +async def test_data_touched_since_delivery_voids_the_guarantee(): + """THE safety property: loading a universe is documented as destructive. If anything has + changed the data since the provider loaded it, every service is re-seeded.""" + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe(), changelog_empty=False) + assert sorted(run.loaded) == sorted(SERVICES) + + +@pytest.mark.asyncio +async def test_an_unreadable_changelog_loads_rather_than_assumes(): + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe(), + changelog_error=RuntimeError("exec failed")) + assert sorted(run.loaded) == sorted(SERVICES) + + +@pytest.mark.asyncio +async def test_container_mode_loads_as_usual(): + """The changelog gate reads the VM's database; a container-mode env has none to read.""" + run = await _load(_env(preloaded={"id": "uni", "version": 39}, mode=SANDBOX_MODE_CONTAINER), _universe()) + assert sorted(run.loaded) == sorted(SERVICES) + + +@pytest.mark.asyncio +async def test_without_the_guarantee_a_snapshot_is_still_baked(): + """The bake is skipped only because nothing was ingested, not for every load.""" + env = _env(preloaded=None) + await _load(env, _universe()) + env._bake_snapshot_after_load.assert_awaited_once() + + +def test_the_guarantee_survives_the_instance_record(): + """The instance store registers a record from its dict form and loads it back, which is + how the env that runs the load comes to hold it.""" + record = DeployedGatewayEnv( + env_id="env-1", env_version=1, sandbox_id="sb-1", gateway_url="https://gw.example", + preloaded_universe={"id": "uni", "version": 39}, + ) + back = DeployedEnv.from_dict(dataclasses.asdict(record)) + assert type(back) is DeployedGatewayEnv + assert back.preloaded_universe == {"id": "uni", "version": 39} + + +def _universe_named(name: str, version: int = 39): + ua = _universe(version) + ua.id = name + return ua + + +@pytest.mark.asyncio +async def test_loading_another_universe_voids_the_guarantee(): + """A load reinstalls each service's changelog triggers, which empties the changelog. A + guarantee for A surviving a load of B would pass that gate on the next request for A, + and skip a load the data needs.""" + env = _env(preloaded={"id": "uni", "version": 39}) + await _load(env, _universe_named("other")) + + run = await _load(env, _universe_named("uni")) + + assert sorted(run.loaded) == sorted(SERVICES) + + +@pytest.mark.asyncio +async def test_the_void_reaches_the_stored_record(): + """So a reattach, which rebuilds the deployment from the record, can't revive it.""" + env = _env(preloaded={"id": "uni", "version": 39}) + + run = await _load(env, _universe_named("other")) + + run.instances.clear_preloaded_universe.assert_called_once_with("inst-1") + assert env._deployed.preloaded_universe is None + + +@pytest.mark.asyncio +async def test_a_skipped_load_keeps_the_guarantee(): + """Nothing replaced the data, so asking for the same universe again skips again.""" + env = _env(preloaded={"id": "uni", "version": 39}) + + run = await _load(env, _universe()) + + run.instances.clear_preloaded_universe.assert_not_called() + assert env._deployed.preloaded_universe == {"id": "uni", "version": 39} + assert (await _load(env, _universe())).loaded == [] + + +@pytest.mark.asyncio +async def test_a_deployment_with_no_guarantee_writes_nothing(): + run = await _load(_env(preloaded=None), _universe()) + + run.instances.clear_preloaded_universe.assert_not_called() + + +@pytest.mark.asyncio +async def test_without_the_flag_a_preloaded_env_still_loads(): + """Off by default, so no existing caller's load changes: an explicit load stays a load + unless the caller asks to skip.""" + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe(), skip_if_preloaded=False) + assert sorted(run.loaded) == sorted(SERVICES) + assert run.result.delivered_preloaded is False + + +@pytest.mark.asyncio +async def test_a_skip_is_reported_on_the_result(): + run = await _load(_env(preloaded={"id": "uni", "version": 39}), _universe()) + assert run.result.delivered_preloaded is True + + +@pytest.mark.asyncio +async def test_loading_a_single_service_voids_the_guarantee(): + """Reloading one service reinstalls its changelog triggers, emptying its part of the + changelog, so a guarantee for A surviving it would skip a load the data needs.""" + env = _env(preloaded={"id": "uni", "version": 39}) + gmail = MagicMock(environment_name="gmail", load_environment_artifact=AsyncMock()) + env.mcp_server_envs = [gmail] + instances = MagicMock() + with patch("agent_env.env.store.get_env_instance_store", return_value=instances): + await env.load_environment_artifact(MagicMock(environment_name="gmail")) + + gmail.load_environment_artifact.assert_awaited_once() + instances.clear_preloaded_universe.assert_called_once_with("inst-1") + assert env._deployed.preloaded_universe is None + + run = await _load(env, _universe()) + assert sorted(run.loaded) == sorted(SERVICES) diff --git a/tst/unit/env/store_test.py b/tst/unit/env/store_test.py index 1a401aa6..7fe6ceff 100644 --- a/tst/unit/env/store_test.py +++ b/tst/unit/env/store_test.py @@ -36,6 +36,7 @@ def test_deployed_env_roundtrip_preserves_all_fields(): created_at_utc=None, expires_at_utc=None, gateway_mode="performance", + preloaded_universe={"id": "uni", "version": 39}, sandbox_ids={ "gateway_server": "sb-gw", "mcp_server": {"slack": "sb-1", "email": "sb-2"}, @@ -232,3 +233,19 @@ def test_instance_ids_do_not_follow_a_reseeded_random(): ids.append(store.create_instance(DeployedEnv(env_id="portsim", env_version=1), ttl_seconds=60).instance_id) assert ids[0] != ids[1] + + +def test_clearing_the_preloaded_universe_survives_a_reload(): + """A reattach rebuilds the deployment from this record; a voided guarantee must stay void.""" + set_document_store(FakeDocumentStore()) + store = EnvInstanceStore() + registered = store.create_instance( + DeployedGatewayEnv(env_id="e", env_version=1, sandbox_id="sb", gateway_url="https://gw.example/", + preloaded_universe={"id": "uni", "version": 39}), + ttl_seconds=60, + ) + assert store.get(registered.instance_id).preloaded_universe == {"id": "uni", "version": 39} + + store.clear_preloaded_universe(registered.instance_id) + + assert store.get(registered.instance_id).preloaded_universe is None diff --git a/tst/unit/task_step/deploy_env_test.py b/tst/unit/task_step/deploy_env_test.py index 119e8798..33657672 100644 --- a/tst/unit/task_step/deploy_env_test.py +++ b/tst/unit/task_step/deploy_env_test.py @@ -433,3 +433,65 @@ def test_deployment_settings_keep_their_positions(): they did: the fifth positional is still ttl_seconds.""" step = DeployEnvTaskStep("t-1.deploy_env", None, "art-1", 2, 60) assert (step.ttl_seconds, step.artifact_id, step.artifact_version) == (60, None, None) + + +def _universe_env(delivers: bool): + """A MultiEnv whose provider does or doesn't promise its deploy step's universe.""" + from types import SimpleNamespace + + from agent_env.env.envs.multi_env import MultiEnv + + env = MagicMock(spec=MultiEnv) + env.deploy = AsyncMock(return_value=_deployed_env(None)) + env._env_provider = SimpleNamespace(delivers_universe=delivers) + env.load_environment_universe_artifact = AsyncMock( + return_value=SimpleNamespace(restored_from_snapshot=False, delivered_preloaded=True) + ) + return env + + +async def _deploy_with_universe(env, **step_fields): + from agent_env.artifact import EnvironmentUniverseArtifact + + universe = MagicMock(spec=EnvironmentUniverseArtifact, id="u", version=3) + step = DeployEnvTaskStep(id="t-1.deploy_env", version=None, env_id="art-1", env_version=2, + **step_fields) + context = TaskStepContext() + with patch("agent_env.env.env.Env") as Env, patch("agent_env.artifact.Artifact.get", + return_value=universe): + Env.get.return_value = env + await step.execute(context) + return universe, context + + +@pytest.mark.asyncio +async def test_a_provider_that_delivers_universes_gets_its_universe_loaded_on_deploy(): + """deploy_env naming a universe is a contract for such a provider: hit or miss, the env + holds it afterwards, so its task needs no load_artifact. The load is skipped when the + env arrived with it, and recorded as load_artifact records it.""" + env = _universe_env(delivers=True) + + universe, context = await _deploy_with_universe(env, artifact_id="u", artifact_version=3) + + env.load_environment_universe_artifact.assert_awaited_once_with(universe, skip_if_preloaded=True) + assert context.metadata["loaded_environment_universes"] == [ + {"id": "u", "version": 3, "restored_from_snapshot": False, "delivered_preloaded": True} + ] + + +@pytest.mark.asyncio +async def test_every_other_provider_leaves_loading_to_load_artifact(): + env = _universe_env(delivers=False) + + await _deploy_with_universe(env, artifact_id="u", artifact_version=3) + + env.load_environment_universe_artifact.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_deploy_naming_no_universe_loads_nothing(): + env = _universe_env(delivers=True) + + await _deploy_with_universe(env) + + env.load_environment_universe_artifact.assert_not_awaited() diff --git a/tst/unit/task_step/test_load_artifact_snapshot_reporting.py b/tst/unit/task_step/test_load_artifact_snapshot_reporting.py index b37e333c..f561a1d0 100644 --- a/tst/unit/task_step/test_load_artifact_snapshot_reporting.py +++ b/tst/unit/task_step/test_load_artifact_snapshot_reporting.py @@ -99,3 +99,25 @@ def test_a_none_result_records_nothing(): context = _context() LoadArtifactTaskStep._record_universe_load(context, SimpleNamespace(id="uni", version=1), None) assert "loaded_environment_universes" not in context.metadata + + +def test_records_a_skip_for_a_preloaded_env(): + context = _context() + LoadArtifactTaskStep._record_universe_load( + context, + SimpleNamespace(id="uni", version=39), + SimpleNamespace(restored_from_snapshot=False, snapshot_db_image_artifact_id=None, + snapshot_baked=None, snapshot_bake_error=None, delivered_preloaded=True), + ) + assert context.metadata["loaded_environment_universes"][0]["delivered_preloaded"] is True + + +def test_a_normal_load_records_no_skip(): + context = _context() + LoadArtifactTaskStep._record_universe_load( + context, + SimpleNamespace(id="uni", version=39), + SimpleNamespace(restored_from_snapshot=False, snapshot_db_image_artifact_id=None, + snapshot_baked=None, snapshot_bake_error=None, delivered_preloaded=False), + ) + assert "delivered_preloaded" not in context.metadata["loaded_environment_universes"][0]