UN-3636 [MISC] Merge overlay manifests via UNSTRACT_RIG_EXTRA_MANIFESTS

backend/utils/tests/test_user_context.py

@@ -0,0 +1,41 @@
+"""Regression tests for UserContext.get_organization.
+
+Pins the no-org short-circuit: the lookup must return None without touching the
+DB, so it stays evaluable on a DB-less/unmigrated setup.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from utils.user_context import UserContext
+
+
+class TestGetOrganizationNoContext:
+    def test_returns_none_without_hitting_db(self):
+        with (
+            patch("utils.user_context.StateStore.get", return_value=None),
+            patch("utils.user_context.Organization.objects.get") as mock_get,
+        ):
+            assert UserContext.get_organization() is None
+            mock_get.assert_not_called()
+
+    def test_empty_identifier_short_circuits(self):
+        with (
+            patch("utils.user_context.StateStore.get", return_value=""),
+            patch("utils.user_context.Organization.objects.get") as mock_get,
+        ):
+            assert UserContext.get_organization() is None
+            mock_get.assert_not_called()
+
+    def test_identifier_present_looks_up_organization(self):
+        """Pins the complement, so inverting the guard can't pass unnoticed."""
+        sentinel = object()
+        with (
+            patch("utils.user_context.StateStore.get", return_value="org-123"),
+            patch(
+                "utils.user_context.Organization.objects.get", return_value=sentinel
+            ) as mock_get,
+        ):
+            assert UserContext.get_organization() is sentinel
+            mock_get.assert_called_once_with(organization_id="org-123")

backend/utils/user_context.py

@@ -18,6 +18,9 @@ def set_organization_identifier(organization_identifier: str) -> None:
     @staticmethod
     def get_organization() -> Organization | None:
         organization_id = StateStore.get(Account.ORGANIZATION_ID)
+        # Skip the query so this stays evaluable on a DB-less/unmigrated setup.
+        if not organization_id:
+            return None
         try:
             organization: Organization = Organization.objects.get(
                 organization_id=organization_id

docker/docker-compose.yaml

@@ -39,7 +39,6 @@ services:
       - ./workflow_data:/data
       - ${TOOL_REGISTRY_CONFIG_SRC_PATH}:/data/tool_registry_config
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-backend
     labels:
       - traefik.enable=true
@@ -60,7 +59,6 @@ services:
       - rabbitmq
       - db
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-metrics
 
     labels:
@@ -83,7 +81,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-ide-callback
       - WORKER_TYPE=ide_callback
       - WORKER_NAME=ide-callback-worker
@@ -107,7 +104,6 @@ services:
     ports:
       - "5555:5555"
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-celery-flower
     volumes:
       - unstract_data:/data
@@ -129,7 +125,6 @@ services:
       - db
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-celery-beat
 
   # Frontend React app
@@ -143,7 +138,6 @@ services:
       - backend
       - reverse-proxy
     environment:
-      - ENVIRONMENT=development
       # Running platform version, surfaced on the profile page. Reuses the same
       # ${VERSION} that tags the image, so it always matches what's deployed
       # (and updates on RC->stable promotion, which only retags — doesn't rebuild).
@@ -219,7 +213,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-api-deployment-v2
       - WORKER_TYPE=api_deployment
       - CELERY_QUEUES_API_DEPLOYMENT=${CELERY_QUEUES_API_DEPLOYMENT:-celery_api_deployments}
@@ -252,7 +245,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-callback-v2
       - WORKER_TYPE=callback
       - WORKER_NAME=callback-worker-v2
@@ -294,7 +286,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-file-processing-v2
       - WORKER_TYPE=file_processing
       - WORKER_MODE=oss
@@ -332,7 +323,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-general-v2
       - WORKER_TYPE=general
       - WORKER_NAME=general-worker-v2
@@ -360,7 +350,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-notification-v2
       - WORKER_TYPE=notification
       - WORKER_NAME=notification-worker-v2
@@ -409,7 +398,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-log-consumer-v2
       - WORKER_TYPE=log_consumer
       - WORKER_NAME=log-consumer-worker-v2
@@ -458,7 +446,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-log-history-scheduler-v2
       # Scheduler interval in seconds
       - LOG_HISTORY_CONSUMER_INTERVAL=${LOG_HISTORY_CONSUMER_INTERVAL:-5}
@@ -483,7 +470,6 @@ services:
       - redis
       - rabbitmq
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-scheduler-v2
       - WORKER_TYPE=scheduler
       - WORKER_NAME=scheduler-worker-v2
@@ -529,7 +515,6 @@ services:
       - rabbitmq
       - platform-service
     environment:
-      - ENVIRONMENT=development
       - APPLICATION_NAME=unstract-worker-executor-v2
       - WORKER_TYPE=executor
       - WORKER_NAME=executor-worker-v2

tests/README.md

@@ -158,7 +158,7 @@ The `platform` pytest fixture in `tests/e2e/conftest.py` reads those env vars; e
 
 Execute-path e2e tests must not call a real provider, so the rig sets `UNSTRACT_LLM_MOCK_RESPONSE` (default `MOCK_LLM_OK`) before boot, for any runtime and treating an exported empty string as unset. The test overlay forwards it into the workers, and `unstract.sdk1.llm` passes it to litellm as `mock_response`: for the non-streaming completion path litellm returns the string verbatim with fixed usage (10 prompt / 20 completion / 30 total), so both the answer and the token counts are exact-assertable. Streaming (`stream_complete`) goes through a different litellm path whose usage differs, so don't assume 10/20/30 there. Sentinels like `litellm.RateLimitError` force error paths. Unset (the production default) the hook is a no-op.
 
-Mocking needs two conditions, not one: `ENVIRONMENT` must also be `test` or `development`. Production sets neither variable, so a stray mock var alone cannot fake completions and their billing — and a refusal is logged rather than silent, since setting the var at all means someone expected mocking. `development` is allowed because that is what base compose sets on the workers that run the injection; the test overlay pins `ENVIRONMENT=test` on those same two workers explicitly, so the tier can't lose its mock to a base-compose edit.
+The variable being set is the only condition, deliberately. A second gate on `ENVIRONMENT` was tried and removed: it did not defend against the case it was written for — a worker env block copied out of this overlay carries both variables together — and compose declared `ENVIRONMENT` on every service anyway, so any deployment derived from it satisfied the gate. Nothing read that variable, so it was dropped everywhere along with the gate. What remains is that the hatch is off unless someone sets it, and that it warns once per process while active. Making fake spend safe to *detect* rather than trying to prevent the config wants the usage record itself to say it was mocked.
 
 A CI/dev override wins (the rig only fills an unset value). Running these tests under the rig **fails** if the var is missing; running **without** the rig just skips the execute-path tests — export it on both sides (your shell and the workers) if you boot the stack yourself.
 

tests/compose/docker-compose.test.yaml

@@ -7,33 +7,23 @@ services:
     # contributor run without setting it.
     image: unstract/backend:${UNSTRACT_TEST_VERSION:-latest}
     environment:
-      - ENVIRONMENT=test
       # Workers ask the backend for this, so this is the value that normally
       # takes effect for fan-out.
       - MAX_PARALLEL_FILE_BATCHES=${MAX_PARALLEL_FILE_BATCHES:-3}
 
   platform-service:
     image: unstract/platform-service:${UNSTRACT_TEST_VERSION:-latest}
-    environment:
-      - ENVIRONMENT=test
 
   runner:
     image: unstract/runner:${UNSTRACT_TEST_VERSION:-latest}
-    environment:
-      - ENVIRONMENT=test
 
-  # Execute-path e2e must never reach a real provider. The delay gives each
-  # mocked completion a known cost so per-file durations stay comparable, and
-  # ENVIRONMENT is the mock's second condition — set here rather than inherited
-  # so the tier can't lose its mock to a base-compose edit.
+  # Execute-path e2e must never reach a real provider.
   worker-executor-v2:
     environment:
-      - ENVIRONMENT=test
       - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-}
 
   worker-file-processing-v2:
     environment:
-      - ENVIRONMENT=test
       - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-}
 
   # Only the worker's fallback if it can't reach the backend; kept in step so it

tests/rig/groups.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import graphlib
+import os
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -26,6 +27,10 @@
 REPO_ROOT = Path(__file__).resolve().parents[2]
 DEFAULT_MANIFEST = REPO_ROOT / "tests" / "groups.yaml"
 
+# os.pathsep-separated overlay manifests, so a downstream repo can contribute
+# groups without editing the base manifest.
+EXTRA_MANIFESTS_ENV = "UNSTRACT_RIG_EXTRA_MANIFESTS"
+
 
 @dataclass(frozen=True)
 class GroupDefinition:
@@ -118,15 +123,20 @@ def expand(self, selected: list[str]) -> list[str]:
 def load_groups(path: Path | None = None) -> GroupManifest:
     """Parse the YAML manifest and validate it."""
     manifest_path = path or DEFAULT_MANIFEST
-    raw = yaml.safe_load(manifest_path.read_text())
-    if not isinstance(raw, dict) or "groups" not in raw:
-        raise ValueError(f"{manifest_path}: expected top-level `groups:` mapping")
+    raw = _load_manifest_dict(manifest_path)
 
     defaults = raw.get("defaults") or {}
     groups: dict[str, GroupDefinition] = {}
     for name, spec in (raw["groups"] or {}).items():
         groups[name] = _build_group(name, spec, defaults)
 
+    # Merge before validation so cross-manifest `depends_on` and the platform
+    # gate are checked over the union. Keyed on `path` being omitted rather than
+    # on its value: an explicit path means "load exactly this", however spelled.
+    if path is None:
+        for extra in _extra_manifest_paths():
+            defaults = _merge_manifest(groups, extra, defaults)
+
     _validate_no_cycles(groups)
     _validate_dep_targets_exist(groups)
     _validate_paths(groups)
@@ -135,6 +145,58 @@ def load_groups(path: Path | None = None) -> GroupManifest:
     return GroupManifest(groups=groups)
 
 
+def _extra_manifest_paths() -> list[Path]:
+    """Overlay manifest paths from ``UNSTRACT_RIG_EXTRA_MANIFESTS``.
+
+    Relative paths resolve against ``REPO_ROOT`` so a downstream repo can point
+    at a manifest it copied into this tree (e.g. ``tests/groups.cloud.yaml``).
+    """
+    raw = os.environ.get(EXTRA_MANIFESTS_ENV, "").strip()
+    if not raw:
+        return []
+    paths: list[Path] = []
+    for entry in filter(None, (e.strip() for e in raw.split(os.pathsep))):
+        p = Path(entry)
+        p = p if p.is_absolute() else REPO_ROOT / p
+        # Name the env var — a bare FileNotFoundError won't say where the path
+        # came from.
+        if not p.is_file():
+            raise ValueError(
+                f"{EXTRA_MANIFESTS_ENV}: {entry!r} is not a file (resolved to {p})"
+            )
+        paths.append(p)
+    return paths
+
+
+def _load_manifest_dict(manifest_path: Path) -> dict[str, Any]:
+    """Parse a manifest YAML, rejecting anything that is not a ``groups:`` mapping."""
+    if not manifest_path.is_file():
+        raise ValueError(f"{manifest_path}: manifest file not found")
+    raw = yaml.safe_load(manifest_path.read_text())
+    if not isinstance(raw, dict) or "groups" not in raw:
+        raise ValueError(f"{manifest_path}: expected top-level `groups:` mapping")
+    return raw
+
+
+def _merge_manifest(
+    groups: dict[str, GroupDefinition], manifest_path: Path, base_defaults: dict[str, Any]
+) -> dict[str, Any]:
+    """Merge an overlay manifest into ``groups`` in place and return the merged
+    defaults, so an overlay can also rename the platform gate. Overlay groups
+    inherit the base ``defaults`` unless they declare their own; a name collision
+    is an error rather than a silent override.
+    """
+    raw = _load_manifest_dict(manifest_path)
+    defaults = {**base_defaults, **(raw.get("defaults") or {})}
+    for name, spec in (raw["groups"] or {}).items():
+        if name in groups:
+            raise ValueError(
+                f"{manifest_path}: group {name!r} already defined in a prior manifest"
+            )
+        groups[name] = _build_group(name, spec, defaults)
+    return defaults
+
+
 def _build_group(
     name: str, spec: dict[str, Any], defaults: dict[str, Any]
 ) -> GroupDefinition:

tests/rig/selection.py

@@ -7,6 +7,9 @@
 The literal ``all`` expands to every group in the manifest. When the result is
 empty, callers should treat that as "do nothing" rather than silently running
 everything — the CLI surfaces a clear error.
+
+With ``--tier``, dep expansion stays inside that tier: tiers run as separate CI
+legs, so a cross-tier ``depends_on`` would re-run the same group in every leg.
 """
 
 from __future__ import annotations
@@ -66,7 +69,16 @@ def resolve(
     if changed_only:
         selected.update(_groups_for_changed_paths(manifest, base=changed_base))
 
-    return manifest.expand(sorted(selected)) if selected else []
+    if not selected:
+        return []
+
+    requested = sorted(selected)
+    expanded = manifest.expand(requested)
+    if tier is not None:
+        # Only dep-expanded groups are dropped; an explicit request survives.
+        keep = set(requested)
+        expanded = [n for n in expanded if n in keep or manifest.get(n).tier == tier]
+    return expanded
 
 
 def _groups_for_changed_paths(manifest: GroupManifest, *, base: str) -> set[str]: