From fc27dfdc7d9be704a1c5bffb9c954caeae9dc886 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 20:37:41 -0700 Subject: [PATCH] test(integration): give each owned proxy a private database and Redis by default (#47213) * test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing. --- tests/integration/_support/actors.py | 4 +- .../integration/_support/bedrock_discovery.py | 3 +- tests/integration/_support/client.py | 28 ++- tests/integration/_support/process.py | 205 ++++++++++++++---- .../integration/_support/prometheus_series.py | 3 +- tests/integration/_support/safeguards_rig.py | 3 +- tests/integration/_support/workers.py | 8 +- .../authorization/test_access_group_budget.py | 2 + .../test_alias_chain_budget_waiver.py | 2 +- .../test_claude_code_gateway_oauth.py | 2 + .../test_cli_sso_login_ui_disabled.py | 2 +- .../authorization/test_end_user_block.py | 2 + .../test_hidden_alias_budget_bypass.py | 2 +- .../test_hidden_alias_budget_bypass_chaos.py | 2 +- ...test_jwt_auto_register_map_existing_key.py | 2 +- .../test_jwt_default_team_provisioning.py | 2 + .../test_jwt_key_mapping_lifecycle.py | 2 +- .../test_jwt_mapped_key_email_backfill.py | 2 + .../test_key_guardrail_opt_out.py | 3 + .../test_key_guardrail_opt_out_chaos.py | 3 + .../test_key_guardrail_opt_out_runtime.py | 3 + .../test_rag_query_vector_store_allowlist.py | 2 + .../test_rag_team_provider_credentials.py | 2 +- .../test_responses_compact_id_ownership.py | 2 + ...st_responses_compact_id_ownership_chaos.py | 2 + .../test_responses_id_ownership.py | 2 + .../authorization/test_sso_blank_identity.py | 2 +- .../test_tool_search_replay_guardrails.py | 2 +- .../caching/test_cache_max_messages.py | 2 + .../caching/test_response_cache_chaos.py | 2 + .../compatibility/_status_code_audit.py | 2 +- .../compatibility/test_codex_model_catalog.py | 2 + .../test_missing_body_param_status.py | 2 + .../test_openai_sdk_client_paths_chaos.py | 2 + .../compatibility/test_persisted_toolsets.py | 2 + .../test_status_code_followups.py | 2 +- .../test_status_code_followups_chaos.py | 2 +- .../configuration/test_codex_catalog_boot.py | 2 + .../configuration/test_effective_settings.py | 2 + tests/integration/conftest.py | 30 ++- .../test_base_model_cross_provider_spend.py | 2 +- ...test_bridged_previous_response_id_spend.py | 2 +- ...est_responses_websocket_guardrail_spend.py | 2 +- .../test_reader_writer_regeneration.py | 2 + .../management/test_batch_file_usage_caps.py | 2 +- .../test_batch_file_usage_caps_chaos.py | 2 +- .../management/test_budget_reset_windows.py | 2 +- .../management/test_budget_updates.py | 2 + ...ude_code_plugin_installation_preference.py | 2 + .../test_credential_federation_values.py | 2 + .../test_fallback_management_public_names.py | 2 +- .../management/test_project_lifecycle.py | 2 + .../management/test_ptu_share_guards.py | 2 +- .../management/test_scim_admin_group.py | 2 + .../management/test_scim_group_bulk_roster.py | 2 + .../test_scim_group_pathless_patch.py | 2 + .../management/test_search_tool_crud.py | 3 + .../management/test_sso_sign_in.py | 2 +- .../test_team_budget_duration_defaults.py | 2 + .../management/test_team_delete_prometheus.py | 2 + .../management/test_tool_policy_user.py | 2 + ...test_vector_store_file_list_managed_ids.py | 2 + .../mcp/test_mcp_accounting_guardrails.py | 2 + .../mcp/test_mcp_agent_365_guardrail.py | 2 + .../mcp/test_mcp_akto_logging_only.py | 2 + .../integration/mcp/test_mcp_llm_endpoints.py | 2 +- tests/integration/mcp/test_mcp_management.py | 2 + tests/integration/mcp/test_mcp_oauth_flows.py | 2 + .../mcp/test_mcp_protocol_errors.py | 2 + tests/integration/mcp/test_mcp_resilience.py | 2 + .../integration/mcp/test_mcp_user_env_vars.py | 2 + .../mcp/test_oauth_configuration.py | 2 + ...chat_bridge_multi_tool_call_blocks_wire.py | 2 + .../anthropic/test_anthropic_gateway_wire.py | 2 + ...hropic_messages_pre_content_retry_chaos.py | 2 +- .../test_anthropic_safeguards_beta_wire.py | 2 + .../test_passthrough_migration_wire.py | 4 + .../test_pdf_document_count_tokens_chaos.py | 2 + .../test_websearch_interception_wire.py | 2 + ...st_anthropic_messages_caller_tools_wire.py | 2 + .../observability/_openinference_support.py | 8 +- .../observability/azure_dispatch_support.py | 3 +- .../observability/test_akto_logging_only.py | 2 + .../test_azure_content_safety_audit.py | 2 + .../test_azure_content_safety_endpoints.py | 2 + .../test_azure_storage_client_ttl.py | 4 + ...ock_converse_stream_post_call_guardrail.py | 2 + ...drock_invoke_stream_post_call_guardrail.py | 2 + .../test_cache_hit_guardrail_metrics.py | 2 + .../test_cache_hit_guardrail_metrics_chaos.py | 3 + .../observability/test_callback_delivery.py | 2 + .../test_content_filter_logging_only.py | 2 + .../test_custom_logger_redaction.py | 2 +- .../observability/test_guardrail_effects.py | 2 + .../test_guardrail_stream_scope.py | 2 + .../test_guardrail_stream_scope_chaos.py | 2 + .../test_guardrail_stream_scope_matrix.py | 2 + .../observability/test_langfuse_delivery.py | 2 + .../observability/test_langtrace_delivery.py | 4 +- .../test_logging_only_scope_chaos.py | 2 + .../test_logging_only_scope_config.py | 2 + .../test_logging_only_scope_runtime.py | 2 + .../observability/test_otel_cache_hit_cost.py | 2 + .../test_otel_conversation_id.py | 2 + .../test_otel_trace_response_header.py | 2 + .../test_passthrough_upstream_error_chaos.py | 2 + ...t_passthrough_upstream_error_visibility.py | 2 + .../test_s3_v2_flush_surfaces.py | 2 + .../test_s3_v2_partition_granularity.py | 2 + .../observability/test_signoz_delivery.py | 2 + .../test_straiker_v3_platform.py | 2 + .../test_stream_error_alert_api_base.py | 2 +- .../observability/test_wingback_guardrail.py | 2 + .../pricing/test_configured_prices.py | 2 + .../pricing/test_price_data_reload_prices.py | 2 + .../providers/_idle_upstream_support.py | 10 +- ...est_anthropic_refusal_stop_details_wire.py | 1 + ...thropic_thinking_signature_logging_wire.py | 2 +- ...pic_tool_search_result_passthrough_wire.py | 2 +- ...bedrock_converse_lookaround_regex_chaos.py | 2 + ...t_bedrock_converse_missing_content_wire.py | 2 + ...drock_converse_stream_event_frames_wire.py | 2 + .../test_bedrock_kimi_k3_cache_point_wire.py | 2 + .../test_bedrock_mantle_claude_chat_chaos.py | 2 + .../test_bedrock_mantle_output_limit_wire.py | 2 + ...che_control_tool_call_marks_owned_proxy.py | 2 +- ...test_cache_control_tool_call_marks_wire.py | 2 + ...t_chatgpt_responses_caller_stream_chaos.py | 2 +- ...st_chatgpt_responses_caller_stream_wire.py | 2 +- .../test_coralbricks_provider_chaos.py | 2 + .../test_coralbricks_provider_env_wire.py | 2 + .../test_databricks_decisions_wire.py | 2 + .../test_decision_model_guardrail.py | 2 + ...est_decisions_box_deployment_flag_chaos.py | 2 + .../providers/test_decisions_chaos.py | 2 + .../test_device_code_login_guard_boot.py | 2 +- .../test_device_code_login_guard_wire.py | 2 +- .../providers/test_fal_ai_passthrough_wire.py | 2 + .../test_gemini_embedding_file_block_chaos.py | 2 + .../test_gemini_long_decode_chaos.py | 2 + .../providers/test_gemini_long_decode_wire.py | 2 + .../test_gemini_thinking_replay_wire.py | 2 + ...gemini_tool_result_part_reference_chaos.py | 2 + ...st_gemini_tts_voice_mapping_owned_proxy.py | 2 +- ...renceapis_context_window_fallbacks_wire.py | 2 +- .../providers/test_malformed_tools_chaos.py | 2 + .../test_ollama_prompt_tools_wire.py | 2 + ...ect_prompt_cache_breakpoint_owned_proxy.py | 2 + .../test_openai_realtime_transcription_url.py | 2 +- .../test_openai_responses_tool_search_wire.py | 2 +- .../test_openai_responses_websocket_wire.py | 2 +- .../test_realtime_unguarded_voice_turns.py | 2 +- ...bridge_message_cache_control_wire_chaos.py | 2 +- ...es_bridge_prompt_cache_breakpoint_chaos.py | 2 + ...idge_prompt_cache_breakpoint_wire_chaos.py | 2 +- ...test_responses_bridge_tool_search_chaos.py | 2 + ...est_rust_anthropic_safeguards_beta_wire.py | 2 +- .../test_sagemaker_embedding_chaos.py | 2 +- .../test_sagemaker_embedding_wire.py | 2 +- .../test_strands_decider_agentcore_wire.py | 2 +- ...vertex_anthropic_structured_output_wire.py | 2 + .../providers/test_vertex_passthrough_wire.py | 2 +- .../test_voyage_mongodb_host_wire.py | 2 +- ...complexity_router_databricks_classifier.py | 2 + .../test_deployment_mode_routing_wire.py | 2 + ...allback_hop_foreign_encrypted_reasoning.py | 2 + .../routing/test_observed_routing.py | 2 + .../routing/test_payment_required_cooldown.py | 2 +- ...df_document_pre_call_checks_owned_proxy.py | 2 + .../routing/test_prompt_cache_pin_ttl.py | 2 +- .../test_prompt_cache_pin_ttl_chaos.py | 2 +- tests/integration/routing/test_ptu_shares.py | 2 +- .../routing/test_ptu_shares_regex_timing.py | 2 +- ..._vllm_passthrough_stream_usage_settings.py | 2 +- tests/integration/security/_sinks.py | 6 +- .../security/test_fine_tuning_managed_ids.py | 2 + .../test_background_interaction_settlement.py | 2 +- .../spend/test_daily_activity_routes.py | 2 + .../spend/test_daily_tag_spend_retention.py | 2 + .../test_disable_end_user_cost_tracking.py | 2 +- .../test_end_user_spend_without_proxy_user.py | 2 + .../test_expired_ui_session_failure_rows.py | 2 +- .../spend/test_key_budget_lockout.py | 2 + .../spend/test_messages_stream_usage_cost.py | 2 + .../test_normalized_error_long_message.py | 2 + .../test_passthrough_budget_reservation.py | 2 + .../spend/test_passthrough_request_tags.py | 2 + .../spend/test_ptu_share_attribution.py | 2 +- .../test_requester_ip_address_logging.py | 2 + .../test_reset_budget_leader_election.py | 2 + .../test_responses_websocket_turn_rows.py | 2 +- .../spend/test_spend_log_caller_session.py | 5 +- .../test_spend_log_tool_payload_content.py | 2 + .../spend/test_spend_log_write_batching.py | 2 + .../spend/test_spend_logs_metadata_fields.py | 2 + .../spend/test_team_member_budget_alerts.py | 2 + .../spend/test_team_member_spend_flush.py | 2 + .../spend/test_user_budget_on_team_keys.py | 2 + .../test_vertex_image_regional_uplift.py | 2 + ...test_chat_stream_guardrail_error_events.py | 1 + .../test_faked_responses_stream_chaos.py | 2 +- ...google_genai_stream_usage_settings_wire.py | 2 + ...t_interception_converted_stream_options.py | 2 +- .../test_noma_v2_streaming_settings.py | 2 + .../test_responses_bridge_stream_chaos.py | 2 +- .../test_responses_bridge_stream_errors.py | 2 +- .../streaming/test_stream_contracts.py | 2 + .../test_chat_completions_reasoning_gemini.py | 2 +- .../cross_model/test_responses_cross_model.py | 2 + .../unit/integration_support/test_process.py | 103 ++++++++- 210 files changed, 676 insertions(+), 144 deletions(-) diff --git a/tests/integration/_support/actors.py b/tests/integration/_support/actors.py index 6a1e53bda31..0d628cc3b88 100644 --- a/tests/integration/_support/actors.py +++ b/tests/integration/_support/actors.py @@ -15,7 +15,7 @@ from pydantic import JsonValue from litellm.proxy.utils import PrismaClient, hash_token from tests.integration._support.client import Gateway, Scenario, string_value from tests.integration._support.prisma_clients import connected_proxy_prisma -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import OwnedState, owned_proxy TEAM_ALPHA: Final = "behavior-pin-team-alpha" TEAM_BETA: Final = "behavior-pin-team-beta" @@ -124,7 +124,7 @@ def owned_settings_client( config: Final = {**base, "litellm_settings": {**base["litellm_settings"], **litellm_settings}} path: Final = directory / "owned-settings.yaml" path.write_text(yaml.safe_dump(config)) - with owned_proxy(client.gateway, directory, {}, config=path) as candidate: + with owned_proxy(client.gateway, directory, {}, config=path, state=OwnedState.SHARED_WITH_WORKER) as candidate: yield GatewayClient(candidate, client.replacements) diff --git a/tests/integration/_support/bedrock_discovery.py b/tests/integration/_support/bedrock_discovery.py index 5b3a091dc17..058e0850068 100644 --- a/tests/integration/_support/bedrock_discovery.py +++ b/tests/integration/_support/bedrock_discovery.py @@ -23,7 +23,7 @@ from integration._support.aws_control_plane import ( verified_scope, ) from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, held, object_value, string_value -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from pydantic import JsonValue @@ -115,6 +115,7 @@ def discovery_proxy( config=discovery_config(parent, directory, check_provider_endpoint=check_provider_endpoint), remove_environment=aws_environment_names(), workers=workers, + state=OwnedState.SHARED_WITH_WORKER, ) as owned: wait_for_every_worker(owned.log, workers) yield owned diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 053318c2451..6f407295259 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -35,9 +35,9 @@ def list_value(value: JsonValue) -> Sequence[JsonValue]: def delete_key_if_present(candidate: Gateway, key: str) -> None: digest: Final = sha256(key.encode()).hexdigest() - if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)): + if candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)): candidate.post("/key/delete", {"keys": [key]}) - assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [] + assert candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [] def eventually( @@ -95,6 +95,10 @@ class Gateway: client: httpx.Client key: str upstream_url: str + database_url: str | None = None + + def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return read_rows(query, parameters, database_url=self.database_url) def request( self, @@ -175,7 +179,7 @@ class Scenario: def delete_team(self, identity: str) -> None: self.gateway.post("/team/delete", {"team_ids": [identity]}) - assert read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == [] + assert self.gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == [] def project(self, team_id: str, **fields: JsonValue) -> str: created: Final = self.gateway.post( @@ -188,7 +192,9 @@ class Scenario: def delete_project(self, identity: str) -> None: response: Final = self.gateway.request("DELETE", "/project/delete", {"project_ids": [identity]}) assert response.status_code == 200, response.text - assert read_rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == [] + assert ( + self.gateway.rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == [] + ) def budget(self, **fields: JsonValue) -> str: created: Final = self.gateway.post("/budget/new", fields) @@ -198,7 +204,7 @@ class Scenario: def delete_budget(self, identity: str) -> None: self.gateway.post("/budget/delete", {"id": identity}) - assert read_rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == [] + assert self.gateway.rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == [] def organization(self, **fields: JsonValue) -> str: created: Final = self.gateway.post( @@ -212,7 +218,9 @@ class Scenario: response: Final = self.gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]}) assert response.status_code == 200, response.text assert ( - read_rows('SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,)) + self.gateway.rows( + 'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,) + ) == [] ) self.delete_budget(budget_id) @@ -236,7 +244,7 @@ class Scenario: def delete_user(self, identity: str) -> None: response: Final = self.gateway.request("POST", "/user/delete", {"user_ids": [identity]}) assert response.status_code == 200 and response.json() == 1, response.text - assert read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == [] + assert self.gateway.rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == [] def member(self, team_id: str, role: str = "user") -> str: """Create an internal user and add them to ``team_id``; deleting the user later removes the membership.""" @@ -247,7 +255,7 @@ class Scenario: def delete_key(self, token: str) -> None: self.gateway.post("/key/delete", {"keys": [token]}) hashed: Final = sha256(token.encode()).hexdigest() - assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == [] + assert self.gateway.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == [] info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"]) assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}" @@ -265,7 +273,9 @@ class Scenario: lambda served: identity not in served, seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, ) - assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] + assert ( + self.gateway.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] + ) def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: name: Final = f"integration-{uuid.uuid4().hex}" diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index d5ce6f916a2..b08037e80db 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -7,9 +7,9 @@ import subprocess import sys import time import uuid -from collections.abc import Generator, Iterable, Iterator, Mapping +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager, suppress +from contextlib import AbstractContextManager, contextmanager, suppress from dataclasses import dataclass from enum import Enum from datetime import UTC, datetime, timedelta @@ -17,7 +17,7 @@ from ipaddress import IPv4Address from itertools import accumulate from pathlib import Path from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, TypeAlias from urllib.parse import urlsplit import httpx @@ -33,6 +33,8 @@ from tests.integration._support.budget import import_seconds, readiness_for, rea from tests.integration._support.client import GATEWAY_LIMITS, Gateway, gateway_from_environment from tests.integration._support.database import scratch_database from tests.integration._support.ports import allocate_port +from tests.integration._support.redis_process import owned_redis +from tests.integration._support.workers import gateway_profile, replica_urls if TYPE_CHECKING: import pytest @@ -43,6 +45,75 @@ LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",) CONFIGURATION: Final = TypeAdapter(dict[str, object]) _DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]]) MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml" +OWNED_STATE_VARIABLE: Final = "INTEGRATION_OWNED_STATE" + + +class OwnedState(Enum): + PRIVATE = "private" + SHARED_WITH_WORKER = "shared_with_worker" + + +@dataclass(frozen=True, slots=True) +class PrivateState: + database_url: str + environment: Mapping[str, str] + + +PrivateStateFactory: TypeAlias = Callable[[Path, Mapping[str, str]], AbstractContextManager[PrivateState]] + + +def boots_private_state( + state: OwnedState | None, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + inherited: Mapping[str, str], +) -> bool: + chosen: Final = state or OwnedState(inherited.get(OWNED_STATE_VARIABLE, OwnedState.PRIVATE.value)) + return chosen is OwnedState.PRIVATE and "DATABASE_URL" not in overrides and "DATABASE_URL" not in remove_environment + + +@contextmanager +def private_state(directory: Path, inherited: Mapping[str, str]) -> Iterator[PrivateState]: + with scratch_database("owned") as database, owned_redis(directory, inherited=inherited) as cache: + replica: Final = replica_urls(database, inherited) if gateway_profile(inherited).replica else None + yield PrivateState( + database, + MappingProxyType( + { + "DATABASE_URL": database if replica is None else replica["INTEGRATION_PROXY_DATABASE_URL"], + **( + {} + if replica is None + else {"DATABASE_URL_READ_REPLICA": replica["INTEGRATION_PROXY_READ_REPLICA_URL"]} + ), + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + } + ), + ) + + +@dataclass(frozen=True, slots=True) +class OwnedEnvironment: + overrides: Mapping[str, str] + database_url: str | None + + +@contextmanager +def owned_environment( + directory: Path, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + state: OwnedState | None, + inherited: Mapping[str, str], + private: PrivateStateFactory = private_state, +) -> Iterator[OwnedEnvironment]: + if not boots_private_state(state, overrides, remove_environment, inherited): + yield OwnedEnvironment(overrides, None) + return + with private(directory, inherited) as booted: + kept: Final = {name: value for name, value in booted.environment.items() if name not in remove_environment} + yield OwnedEnvironment(MappingProxyType({**kept, **overrides}), booted.database_url) def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]: @@ -368,6 +439,7 @@ def owned_proxy( database_setup: tuple[str, ...] = DB_PUSH, coverage_file: Path | None = None, inherited: Mapping[str, str] = os.environ, + state: OwnedState | None = None, ) -> Iterator[Gateway]: with owned_proxy_process( gateway, @@ -379,20 +451,21 @@ def owned_proxy( database_setup=database_setup, coverage_file=coverage_file, inherited=inherited, + state=state, ) as owned: yield owned.gateway @contextmanager def isolated_gateway(directory: Path, environment: "pytest.MonkeyPatch") -> Iterator[Gateway]: - with gateway_from_environment() as shared, scratch_database() as database: - environment.setenv("DATABASE_URL", database) - environment.setenv("INTEGRATION_PROXY_DATABASE_URL", database) + with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, state=OwnedState.PRIVATE) as owned: + assert owned.database_url is not None + environment.setenv("DATABASE_URL", owned.database_url) + environment.setenv("INTEGRATION_PROXY_DATABASE_URL", owned.database_url) environment.delenv("INTEGRATION_PROXY_READ_REPLICA_URL", raising=False) environment.delenv("DATABASE_URL_READ_REPLICA", raising=False) - with owned_proxy(shared, directory, {}) as owned: - owned.post("/config/update", {"router_settings": {"num_retries": 0}}) - yield owned + owned.post("/config/update", {"router_settings": {"num_retries": 0}}) + yield owned def _stop(process: subprocess.Popen[bytes]) -> None: @@ -634,6 +707,32 @@ def _proxy_root() -> Path: return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +def _configuration_copy(config: Path, directory: Path, prefix: str) -> Path: + checkouts: Final = (_harness_root(), _proxy_root().resolve()) + inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts) + target: Final = directory if inside_checkout else config.parent + target.mkdir(parents=True, exist_ok=True) + return target / f"{prefix}-{uuid.uuid4().hex}.yaml" + + +def private_configuration(config: Path, directory: Path) -> Path: + loaded: Final[object] = yaml.safe_load(config.read_text()) + if not isinstance(loaded, dict): + return config + settings: Final = CONFIGURATION.validate_python(loaded) + router: Final = settings.get("router_settings") or {} + if not isinstance(router, dict) or "num_retries" in router: + return config + destination: Final = _configuration_copy(config, directory, "private-config") + destination.write_text( + yaml.safe_dump( + {**settings, "router_settings": {**CONFIGURATION.validate_python(router), "num_retries": 0}}, + sort_keys=False, + ) + ) + return destination + + def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path: if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1": return config @@ -644,11 +743,7 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) - checkouts: Final = (_harness_root(), _proxy_root().resolve()) - inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts) - target: Final = directory if inside_checkout else config.parent - target.mkdir(parents=True, exist_ok=True) - destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml" + destination: Final = _configuration_copy(config, directory, "worker-config") loaded: Final[object] = yaml.safe_load(contents) if not isinstance(loaded, dict): destination.write_text(contents) @@ -753,42 +848,58 @@ def owned_proxy_process( healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT, inherited: Mapping[str, str] = os.environ, client_timeout: float = 15, + state: OwnedState | None = None, ) -> Iterator[OwnedProxy]: assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`" root: Final = _proxy_root() - environment: Final = _proxy_environment( - gateway, - {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, - remove_environment, - inherited, - ) - output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) - output.mkdir(parents=True, exist_ok=True) - command: Final = ( - sys.executable, - "-m", - *(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()), - "integration._support.proxy", - "--config", - str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)), - "--host", - "127.0.0.1", - "--num_workers", - str(workers), - "--timeout_worker_healthcheck", - str(healthcheck_seconds(healthcheck)), - *database_setup, - *extra_arguments, - ) - launch: Final = _launch_until_ready(command, root, environment, output, workers) - process: Final = launch.process - try: - with httpx.Client( - base_url=f"http://127.0.0.1:{launch.port}", timeout=client_timeout, trust_env=False, limits=GATEWAY_LIMITS - ) as client: - yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log, healthcheck) - finally: - _stop_launch(launch) + with owned_environment(directory, overrides, remove_environment, state, inherited) as owned: + configured: Final = worker_configuration( + config or root / "tests/integration/proxy_config.yaml", directory, inherited + ) + owned_configuration: Final = ( + configured if owned.database_url is None else private_configuration(configured, directory) + ) + environment: Final = _proxy_environment( + gateway, + {**owned.overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, + remove_environment, + inherited, + ) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + command: Final = ( + sys.executable, + "-m", + *(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()), + "integration._support.proxy", + "--config", + str(owned_configuration), + "--host", + "127.0.0.1", + "--num_workers", + str(workers), + "--timeout_worker_healthcheck", + str(healthcheck_seconds(healthcheck)), + *database_setup, + *extra_arguments, + ) + launch: Final = _launch_until_ready(command, root, environment, output, workers) + process: Final = launch.process + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{launch.port}", + timeout=client_timeout, + trust_env=False, + limits=GATEWAY_LIMITS, + ) as client: + yield OwnedProxy( + Gateway(client, gateway.key, gateway.upstream_url, owned.database_url), + process, + launch.log, + healthcheck, + ) + finally: + _stop_launch(launch) @contextmanager diff --git a/tests/integration/_support/prometheus_series.py b/tests/integration/_support/prometheus_series.py index 3a406a97d11..6f54f2a8f7f 100644 --- a/tests/integration/_support/prometheus_series.py +++ b/tests/integration/_support/prometheus_series.py @@ -19,7 +19,7 @@ import httpx import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process from integration._support.responses_vendor import ResponsesVendor, same_response from integration._support.wire import Reply, Request, Wire, wire_server from prometheus_client.mmap_dict import MmapedDict @@ -207,6 +207,7 @@ def series_cap_rig( config=config, remove_environment=() if shared_samples else ("PROMETHEUS_MULTIPROC_DIR",), workers=workers, + state=OwnedState.SHARED_WITH_WORKER, ) ) scenario: Final = stack.enter_context(owned.gateway.scenario()) diff --git a/tests/integration/_support/safeguards_rig.py b/tests/integration/_support/safeguards_rig.py index 8ba6f14f71c..f0acf956e5e 100644 --- a/tests/integration/_support/safeguards_rig.py +++ b/tests/integration/_support/safeguards_rig.py @@ -30,7 +30,7 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -313,6 +313,7 @@ def safeguards_rig( config=path, remove_environment=remove_environment, workers=workers, + state=OwnedState.SHARED_WITH_WORKER, ) as owned: yield Rig(wire, scripts, owned) diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index e75b0b3c6bb..f1958ef6990 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -86,9 +86,9 @@ def worker_endpoints() -> Mapping[str, str]: ) -def _replica_urls(database: str) -> Mapping[str, str]: - writer: Final = urlsplit(os.environ["INTEGRATION_PROXY_DATABASE_URL"]) - reader: Final = urlsplit(os.environ["INTEGRATION_PROXY_READ_REPLICA_URL"]) +def replica_urls(database: str, inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]: + writer: Final = urlsplit(inherited["INTEGRATION_PROXY_DATABASE_URL"]) + reader: Final = urlsplit(inherited["INTEGRATION_PROXY_READ_REPLICA_URL"]) database_path: Final = urlsplit(database).path assert writer.username is not None and reader.username is not None with psycopg.connect(database, autocommit=True) as admin: @@ -164,7 +164,7 @@ def worker_services( "INTEGRATION_PROXY_READ_REPLICA_URL": "", "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), - **(_replica_urls(database) if replica else {}), + **(replica_urls(database, inherited) if replica else {}), } ) diff --git a/tests/integration/authorization/test_access_group_budget.py b/tests/integration/authorization/test_access_group_budget.py index 0db05a649e3..71302be2f37 100644 --- a/tests/integration/authorization/test_access_group_budget.py +++ b/tests/integration/authorization/test_access_group_budget.py @@ -22,6 +22,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _INPUT_RATE: Final = 0.000001 _OUTPUT_RATE: Final = 0.00001 _MAX_BUDGET: Final = 0.0001 diff --git a/tests/integration/authorization/test_alias_chain_budget_waiver.py b/tests/integration/authorization/test_alias_chain_budget_waiver.py index 0d64415d445..0908fa0691c 100644 --- a/tests/integration/authorization/test_alias_chain_budget_waiver.py +++ b/tests/integration/authorization/test_alias_chain_budget_waiver.py @@ -62,7 +62,7 @@ from integration.authorization._hidden_alias_budget import ( ) from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] PAID_INPUT_COST_PER_TOKEN: Final = 0.001 PAID_OUTPUT_COST_PER_TOKEN: Final = 0.002 diff --git a/tests/integration/authorization/test_claude_code_gateway_oauth.py b/tests/integration/authorization/test_claude_code_gateway_oauth.py index acab042bfc9..384bf3a2918 100644 --- a/tests/integration/authorization/test_claude_code_gateway_oauth.py +++ b/tests/integration/authorization/test_claude_code_gateway_oauth.py @@ -23,6 +23,8 @@ from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _CLIENT_ID: Final = "scripted-cli-client" _CLIENT_SECRET: Final = "scripted-cli-secret" _AUTHORIZATION_CODE: Final = "scripted-auth-code" diff --git a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py index 315a0b20334..a53fbfacf4c 100644 --- a/tests/integration/authorization/test_cli_sso_login_ui_disabled.py +++ b/tests/integration/authorization/test_cli_sso_login_ui_disabled.py @@ -34,7 +34,7 @@ from tests.integration._support.process import OwnedProxy, group_members, owned_ from tests.integration._support.provider import SharedProvider from tests.integration._support.wire import Reply, Request, Wire, wire_server -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] CLI_SOURCE: Final = "litellm-cli" CLI_STATE_PREFIX: Final = "litellm-session-token" diff --git a/tests/integration/authorization/test_end_user_block.py b/tests/integration/authorization/test_end_user_block.py index 8b18bddfd00..fd283380d8f 100644 --- a/tests/integration/authorization/test_end_user_block.py +++ b/tests/integration/authorization/test_end_user_block.py @@ -12,6 +12,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from pydantic import BaseModel, JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + class _BlockedCustomerInfo(BaseModel): user_id: str diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass.py b/tests/integration/authorization/test_hidden_alias_budget_bypass.py index 493db4c3e62..51b36c145d3 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass.py @@ -62,7 +62,7 @@ from integration.authorization._hidden_alias_budget import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] # $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-10-14 in 900a57f5b1f): # Lemonade is a local inference server, so there is no vendor price for the row to drift from. A row that gains a diff --git a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py index 0a9b1cb3c4e..3e6acfc72de 100644 --- a/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py +++ b/tests/integration/authorization/test_hidden_alias_budget_bypass_chaos.py @@ -32,7 +32,7 @@ from integration.authorization._hidden_alias_budget import ( upstream_requests, ) -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _WAVE: Final = 8 diff --git a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py index 5215054a364..62eda977438 100644 --- a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py +++ b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py @@ -24,7 +24,7 @@ from tests.integration._support.wire import Reply, Request, wire_server KEY_ID: Final = "integration-jwt-map-existing-key" MAPPING_INSERT: Final = b'INSERT INTO "public"."LiteLLM_JWTKeyMapping"' -pytestmark = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] def _hash(key: str) -> str: diff --git a/tests/integration/authorization/test_jwt_default_team_provisioning.py b/tests/integration/authorization/test_jwt_default_team_provisioning.py index 04715d128b0..a8c00d50ee9 100644 --- a/tests/integration/authorization/test_jwt_default_team_provisioning.py +++ b/tests/integration/authorization/test_jwt_default_team_provisioning.py @@ -14,6 +14,8 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + KEY_ID: Final = "integration-jwt-signing-key" TEAM_BUDGET: Final = 25.0 diff --git a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py index cc1d5c35788..715f3731175 100644 --- a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py +++ b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py @@ -36,7 +36,7 @@ _KEY_ID: Final = "integration-signing-key" _UPSTREAM_KEY: Final = "synthetic-openai-key" _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue]) -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = [pytest.mark.timeout(180), pytest.mark.shares_worker_state] class _Usage(BaseModel): diff --git a/tests/integration/authorization/test_jwt_mapped_key_email_backfill.py b/tests/integration/authorization/test_jwt_mapped_key_email_backfill.py index 47c8ad5f115..b676ddad494 100644 --- a/tests/integration/authorization/test_jwt_mapped_key_email_backfill.py +++ b/tests/integration/authorization/test_jwt_mapped_key_email_backfill.py @@ -14,6 +14,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from jwt.algorithms import RSAAlgorithm +pytestmark: Final = pytest.mark.shares_worker_state + AUDIENCE: Final = "litellm-integration" KEY_ID: Final = "integration-signing-key" CLIENT_CLAIM: Final = "client_id" diff --git a/tests/integration/authorization/test_key_guardrail_opt_out.py b/tests/integration/authorization/test_key_guardrail_opt_out.py index 6591b0b3993..0f47ea269f9 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Final import httpx +import pytest import yaml from integration._support.client import Gateway, Scenario, object_value, string_value @@ -16,6 +17,8 @@ from integration.authorization._guardrail_opt_out import ( stored_metadata, ) +pytestmark: Final = pytest.mark.shares_worker_state + _KEY_ROUTES: Final = ["/key/generate", "/key/update", "/key/regenerate", "/v1/chat/completions"] diff --git a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py index e136691aa2e..3ceae540ba9 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py @@ -9,6 +9,7 @@ from queue import SimpleQueue from typing import Final import httpx +import pytest from integration._support.client import Gateway, eventually, string_value from integration._support.database import read_rows @@ -24,6 +25,8 @@ from integration.authorization._guardrail_opt_out import ( upstream_observations, ) +pytestmark: Final = pytest.mark.shares_worker_state + class _GuardrailSink: """Test-owned guardrail endpoint that can be stopped and restarted on the same port.""" diff --git a/tests/integration/authorization/test_key_guardrail_opt_out_runtime.py b/tests/integration/authorization/test_key_guardrail_opt_out_runtime.py index dc549d3be8d..2b6f1860ea9 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out_runtime.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out_runtime.py @@ -5,6 +5,7 @@ from pathlib import Path from typing import Final import httpx +import pytest from anthropic import Anthropic from openai import AsyncOpenAI, OpenAI @@ -20,6 +21,8 @@ from integration.authorization._guardrail_opt_out import ( upstream_hits, ) +pytestmark: Final = pytest.mark.shares_worker_state + def _wire_hits(wire: Wire, marker: str) -> int: return sum(1 for request in wire.drain() if marker.encode() in request.body) diff --git a/tests/integration/authorization/test_rag_query_vector_store_allowlist.py b/tests/integration/authorization/test_rag_query_vector_store_allowlist.py index f4b5063441f..c1777c66c39 100644 --- a/tests/integration/authorization/test_rag_query_vector_store_allowlist.py +++ b/tests/integration/authorization/test_rag_query_vector_store_allowlist.py @@ -25,6 +25,8 @@ from openai.types.chat import ChatCompletion from pydantic import JsonValue, TypeAdapter from redis import Redis +pytestmark: Final = pytest.mark.shares_worker_state + CONFIG_STORE_ID: Final = "vs_integration_config_store" PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" REMOVE_OPENAI_API_BASE: Final = ("OPENAI_API_BASE",) diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py index 8d556b89107..61bab74ecaf 100644 --- a/tests/integration/authorization/test_rag_team_provider_credentials.py +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -89,7 +89,7 @@ AWS_AMBIENT_CREDENTIALS: Final = ( "AWS_WEB_IDENTITY_TOKEN_FILE", ) -pytestmark: Final = pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS) +pytestmark: Final = [pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS), pytest.mark.shares_worker_state] Signature = tuple[str, str, str, str | None] diff --git a/tests/integration/authorization/test_responses_compact_id_ownership.py b/tests/integration/authorization/test_responses_compact_id_ownership.py index e57118da129..b5a6215dd15 100644 --- a/tests/integration/authorization/test_responses_compact_id_ownership.py +++ b/tests/integration/authorization/test_responses_compact_id_ownership.py @@ -50,6 +50,8 @@ from integration.authorization._responses_compact import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _CLIENTS: Final = ("sync", "async", "httpx") diff --git a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py index d6ff66dc7d1..00327c955ac 100644 --- a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py +++ b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py @@ -40,6 +40,8 @@ from integration.authorization._responses_compact import ( vendor_id, ) +pytestmark: Final = pytest.mark.shares_worker_state + _Kind = Literal["owner", "other_user", "other_team", "unmanaged"] _KINDS: Final[tuple[_Kind, ...]] = ("owner", "other_user", "other_team", "unmanaged") _REFUSING: Final[tuple[_Kind, ...]] = ("other_user", "other_team", "unmanaged") diff --git a/tests/integration/authorization/test_responses_id_ownership.py b/tests/integration/authorization/test_responses_id_ownership.py index 9a88349cca0..2ab26e75f50 100644 --- a/tests/integration/authorization/test_responses_id_ownership.py +++ b/tests/integration/authorization/test_responses_id_ownership.py @@ -19,6 +19,8 @@ from integration._support.worker_readiness import wait_for_worker_readiness from openai.types.responses import ResponseInputParam from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL: Final = "gpt-5" _API_KEY: Final = "synthetic-responses-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/authorization/test_sso_blank_identity.py b/tests/integration/authorization/test_sso_blank_identity.py index 6d9b3b50cc0..4383240f0ba 100644 --- a/tests/integration/authorization/test_sso_blank_identity.py +++ b/tests/integration/authorization/test_sso_blank_identity.py @@ -30,7 +30,7 @@ from tests.integration._support.database_relay import dropped_connection_relay from tests.integration._support.process import OwnedProxy, owned_proxy_process from tests.integration._support.wire import Reply, Request, Wire, wire_server -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] CLIENT_ID: Final = "integration-entra-client" CLIENT_SECRET: Final = "integration-entra-secret" diff --git a/tests/integration/authorization/test_tool_search_replay_guardrails.py b/tests/integration/authorization/test_tool_search_replay_guardrails.py index 2f0ca383201..5b9e29a6d9e 100644 --- a/tests/integration/authorization/test_tool_search_replay_guardrails.py +++ b/tests/integration/authorization/test_tool_search_replay_guardrails.py @@ -18,7 +18,7 @@ from pydantic import JsonValue from tests.integration.caching.response_cache_case import await_entry, redis_store -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _POLICY: Final = "integration-tool-policy" _BLOCK: Final = "tool-permission-block" diff --git a/tests/integration/caching/test_cache_max_messages.py b/tests/integration/caching/test_cache_max_messages.py index 2079761d489..32c71e887fe 100644 --- a/tests/integration/caching/test_cache_max_messages.py +++ b/tests/integration/caching/test_cache_max_messages.py @@ -19,6 +19,8 @@ from pydantic import JsonValue, TypeAdapter from redis import Redis from redis.exceptions import ResponseError +pytestmark: Final = pytest.mark.shares_worker_state + _Turns = Callable[[str], tuple[list[JsonValue], list[JsonValue]]] _INDEXES: Final = TypeAdapter(tuple[bytes, ...]) _COUNT: Final = TypeAdapter(tuple[int]) diff --git a/tests/integration/caching/test_response_cache_chaos.py b/tests/integration/caching/test_response_cache_chaos.py index 17637c088a0..7f330658a65 100644 --- a/tests/integration/caching/test_response_cache_chaos.py +++ b/tests/integration/caching/test_response_cache_chaos.py @@ -61,6 +61,8 @@ from tests.integration.caching.response_cache_case import ( slowed, ) +pytestmark: Final = pytest.mark.shares_worker_state + RecordProperty = Callable[[str, object], None] WORKERS: Final = 2 diff --git a/tests/integration/compatibility/_status_code_audit.py b/tests/integration/compatibility/_status_code_audit.py index 59513891ec0..4ba1ee6b102 100644 --- a/tests/integration/compatibility/_status_code_audit.py +++ b/tests/integration/compatibility/_status_code_audit.py @@ -248,7 +248,7 @@ def owned_gateway( gateway_from_environment() as gateway, owned_proxy_process(gateway, directory, environment or {}, config=path, workers=workers) as owned, ): - yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + yield owned.gateway def stream_finished(endpoint: str, lines: Sequence[str]) -> bool: diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index f1eea96a268..a3ba032e612 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -32,6 +32,8 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + T = TypeVar("T") CLIENT_VERSION: Final = "0.159.3" diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index dd98ff6dea1..673fb0acc51 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -43,6 +43,8 @@ from tests.integration.cost_calculation.cost_tracking_case import ( SseResponse, ) +pytestmark: Final = pytest.mark.shares_worker_state + _Route = tuple[str, str, tuple[str, ...], dict[str, JsonValue], str] _EVAL_DATA_SOURCE_CONFIG: Final[dict[str, JsonValue]] = { "type": "custom", diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py index e9ee4a323cd..9710d8a1b5c 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py @@ -26,6 +26,8 @@ from pydantic import JsonValue from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse +pytestmark: Final = pytest.mark.shares_worker_state + _Kind: TypeAlias = Literal["chat", "chat_stream", "messages", "responses", "completions", "moderations"] _ALL_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream", "messages", "responses", "completions", "moderations") diff --git a/tests/integration/compatibility/test_persisted_toolsets.py b/tests/integration/compatibility/test_persisted_toolsets.py index 80f059b732d..156b28c5cc6 100644 --- a/tests/integration/compatibility/test_persisted_toolsets.py +++ b/tests/integration/compatibility/test_persisted_toolsets.py @@ -12,6 +12,8 @@ from integration._support.database import read_rows from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names from integration._support.process import owned_proxy +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers("other.compatibility.mcp.persisted_tool_names_survive_candidate_startup") def test_existing_toolset_format_loads_before_start_and_keeps_sibling_denied(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/integration/compatibility/test_status_code_followups.py b/tests/integration/compatibility/test_status_code_followups.py index 466eeeddde1..662f37aed1a 100644 --- a/tests/integration/compatibility/test_status_code_followups.py +++ b/tests/integration/compatibility/test_status_code_followups.py @@ -207,7 +207,7 @@ def router_default_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[G }, ) with owned_proxy_process(gateway, directory, {}, config=config) as owned: - yield Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + yield owned.gateway def test_ocr_multipart_without_file_returns_400(gateway: Gateway) -> None: diff --git a/tests/integration/compatibility/test_status_code_followups_chaos.py b/tests/integration/compatibility/test_status_code_followups_chaos.py index 1cb7f45e681..f1ea9e60330 100644 --- a/tests/integration/compatibility/test_status_code_followups_chaos.py +++ b/tests/integration/compatibility/test_status_code_followups_chaos.py @@ -41,7 +41,7 @@ from pydantic import JsonValue from tests.integration.cost_calculation.cost_tracking_case import StoredResponse -pytestmark: Final = pytest.mark.timeout(480) +pytestmark: Final = [pytest.mark.timeout(480), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _MARKER_HEADER: Final = "x-litellm-spend-logs-metadata" diff --git a/tests/integration/configuration/test_codex_catalog_boot.py b/tests/integration/configuration/test_codex_catalog_boot.py index 4868e43b855..1e691f0caaf 100644 --- a/tests/integration/configuration/test_codex_catalog_boot.py +++ b/tests/integration/configuration/test_codex_catalog_boot.py @@ -16,6 +16,8 @@ from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from tests.integration._support.process import owned_proxy, owned_proxy_process +pytestmark: Final = pytest.mark.shares_worker_state + CLIENT_VERSION: Final = "0.159.3" FRESH_CONNECTION: Final = {"Connection": "close"} CONSISTENT_READS: Final = 8 diff --git a/tests/integration/configuration/test_effective_settings.py b/tests/integration/configuration/test_effective_settings.py index 014cfd3a866..029a5c7220e 100644 --- a/tests/integration/configuration/test_effective_settings.py +++ b/tests/integration/configuration/test_effective_settings.py @@ -23,6 +23,8 @@ from tests.integration._support.database import read_rows, scratch_database from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 92fed7200e1..672ab7b6c6a 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -5,7 +5,7 @@ import json import os import re from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from copy import copy from functools import partial from importlib.metadata import version @@ -23,6 +23,7 @@ from tests.integration._support.client import Gateway, Scenario, eventually, gat from tests.integration._support.generation import LIFECYCLE_SETTINGS from tests.integration._support.manifest import OWNED_DIRECTORIES from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups +from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin from tests.integration._support.sdk_models import build_sdk_models @@ -69,6 +70,7 @@ def pytest_addoption(parser: pytest.Parser) -> None: def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "integration: owned real-service integration contracts") config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced") + config.addinivalue_line("markers", "shares_worker_state: owned proxies here share the worker DB and Redis") config.stash[REPORTS] = [] config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files"))) config.pluginmanager.register(IntegrationReportPlugin(config)) @@ -295,6 +297,32 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: session.exitstatus = pytest.ExitCode.OK +@contextmanager +def _owned_state(item: pytest.Item) -> Iterator[None]: + with pytest.MonkeyPatch.context() as environment: + if item.get_closest_marker("shares_worker_state") is not None: + environment.setenv(OWNED_STATE_VARIABLE, OwnedState.SHARED_WITH_WORKER.value) + yield + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_setup(item: pytest.Item) -> Generator[None, None, None]: + with _owned_state(item): + return (yield) + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_call(item: pytest.Item) -> Generator[None, None, None]: + with _owned_state(item): + return (yield) + + +@pytest.hookimpl(wrapper=True) +def pytest_runtest_teardown(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, None, None]: + with _owned_state(item): + return (yield) + + @pytest.fixture def gateway() -> Iterator[Gateway]: with gateway_from_environment() as value: diff --git a/tests/integration/cost_calculation/test_base_model_cross_provider_spend.py b/tests/integration/cost_calculation/test_base_model_cross_provider_spend.py index 02d8ca9cc73..c39ac5345c4 100644 --- a/tests/integration/cost_calculation/test_base_model_cross_provider_spend.py +++ b/tests/integration/cost_calculation/test_base_model_cross_provider_spend.py @@ -16,7 +16,7 @@ from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _MODEL: Final = "cc-subscription-sol" _INPUT_TOKENS: Final = 1000 diff --git a/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py b/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py index 2a8c63fff3d..fcedfcf38ec 100644 --- a/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py +++ b/tests/integration/cost_calculation/test_bridged_previous_response_id_spend.py @@ -17,7 +17,7 @@ from integration._support.workers import cost_map_environment from integration.cost_calculation.cost_tracking_case import COST_MAP from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) diff --git a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py index 8af70d39186..6377be6d095 100644 --- a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py +++ b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py @@ -21,7 +21,7 @@ from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create from integration._support.tls import write_self_signed_cert from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _INPUT_COST_PER_TOKEN: Final = 0.002 _OUTPUT_COST_PER_TOKEN: Final = 0.004 diff --git a/tests/integration/database/test_reader_writer_regeneration.py b/tests/integration/database/test_reader_writer_regeneration.py index 62c152167ab..fc5dbe87e75 100644 --- a/tests/integration/database/test_reader_writer_regeneration.py +++ b/tests/integration/database/test_reader_writer_regeneration.py @@ -17,6 +17,8 @@ from integration._support.wire import Reply, Request, wire_server from psycopg import sql from pydantic import BaseModel, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants") def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/integration/management/test_batch_file_usage_caps.py b/tests/integration/management/test_batch_file_usage_caps.py index 683425f3b1d..2d8186d87ae 100644 --- a/tests/integration/management/test_batch_file_usage_caps.py +++ b/tests/integration/management/test_batch_file_usage_caps.py @@ -59,7 +59,7 @@ from integration.management._batch_file_caps import ( from pydantic import JsonValue from redis import Redis -pytestmark = pytest.mark.timeout(600) +pytestmark: Final = [pytest.mark.timeout(600), pytest.mark.shares_worker_state] YAML_RECORDS: Final = 3 YAML_UPLOADS: Final = 4 diff --git a/tests/integration/management/test_batch_file_usage_caps_chaos.py b/tests/integration/management/test_batch_file_usage_caps_chaos.py index 5f31fa94230..36730278948 100644 --- a/tests/integration/management/test_batch_file_usage_caps_chaos.py +++ b/tests/integration/management/test_batch_file_usage_caps_chaos.py @@ -46,7 +46,7 @@ from integration.management._batch_file_caps import ( ) from pydantic import JsonValue -pytestmark = pytest.mark.timeout(600) +pytestmark: Final = [pytest.mark.timeout(600), pytest.mark.shares_worker_state] WORKERS: Final = 2 UPLOAD_CAP: Final = 5 diff --git a/tests/integration/management/test_budget_reset_windows.py b/tests/integration/management/test_budget_reset_windows.py index 34a0823cea2..999e42bf13d 100644 --- a/tests/integration/management/test_budget_reset_windows.py +++ b/tests/integration/management/test_budget_reset_windows.py @@ -49,7 +49,7 @@ from integration.cost_calculation.cost_tracking_case import JsonResponse, SseRes from pydantic import JsonValue, TypeAdapter from redis import Redis -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] DAY: Final = 86400 HOUR: Final = 3600 diff --git a/tests/integration/management/test_budget_updates.py b/tests/integration/management/test_budget_updates.py index 7a79e35d2af..6939b447866 100644 --- a/tests/integration/management/test_budget_updates.py +++ b/tests/integration/management/test_budget_updates.py @@ -9,6 +9,8 @@ from tests.integration._support.client import Gateway, string_value from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy +pytestmark: Final = pytest.mark.shares_worker_state + def _persisted_reset_at(budget_id: str) -> datetime: rows: Final = read_rows( diff --git a/tests/integration/management/test_claude_code_plugin_installation_preference.py b/tests/integration/management/test_claude_code_plugin_installation_preference.py index b8afd3ed65d..aaf6f02ff4e 100644 --- a/tests/integration/management/test_claude_code_plugin_installation_preference.py +++ b/tests/integration/management/test_claude_code_plugin_installation_preference.py @@ -14,6 +14,8 @@ from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, ev from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import OwnedProxy, owned_proxy_process +pytestmark: Final = pytest.mark.shares_worker_state + PREFERENCES: Final = ("available", "auto_install", "required") INVALID_PREFERENCES: Final[tuple[JsonValue, ...]] = ( "", diff --git a/tests/integration/management/test_credential_federation_values.py b/tests/integration/management/test_credential_federation_values.py index 99be9f0cb8e..2ca9c0cb47c 100644 --- a/tests/integration/management/test_credential_federation_values.py +++ b/tests/integration/management/test_credential_federation_values.py @@ -48,6 +48,8 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from tests.integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + T = TypeVar("T") diff --git a/tests/integration/management/test_fallback_management_public_names.py b/tests/integration/management/test_fallback_management_public_names.py index 75dad3ae428..a5b311e38eb 100644 --- a/tests/integration/management/test_fallback_management_public_names.py +++ b/tests/integration/management/test_fallback_management_public_names.py @@ -29,7 +29,7 @@ from tests.integration._support.process import ( ) from tests.integration._support.redis_process import owned_redis -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] PROVIDER_KEY: Final = "integration-provider-key" EVICTION_WARNING: Final = "config cache eviction of router_settings failed" diff --git a/tests/integration/management/test_project_lifecycle.py b/tests/integration/management/test_project_lifecycle.py index d503daf6a0f..54989c67bea 100644 --- a/tests/integration/management/test_project_lifecycle.py +++ b/tests/integration/management/test_project_lifecycle.py @@ -31,6 +31,8 @@ from pydantic import JsonValue from litellm.models.user import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +pytestmark: Final = pytest.mark.shares_worker_state + _OWNED_PROXY_SALT_KEY: Final = "sk-integration-salt" _OWNERSHIP_POOL_LIMIT: Final = 4 _CONNECTION_BUDGET: Final = re.compile( diff --git a/tests/integration/management/test_ptu_share_guards.py b/tests/integration/management/test_ptu_share_guards.py index 580750bcdef..1b804d1ae7e 100644 --- a/tests/integration/management/test_ptu_share_guards.py +++ b/tests/integration/management/test_ptu_share_guards.py @@ -17,7 +17,7 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 240) +pytestmark: Final = [pytest.mark.timeout(4 * graceful_stop_seconds() + 240), pytest.mark.shares_worker_state] _RELOAD_SECONDS: Final = 5 _PROPAGATION_SECONDS: Final = 6 * _RELOAD_SECONDS diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index b7bf3cfe968..3aeed45d186 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -13,6 +13,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy_process from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + SCIM_HEADERS: Final = {"Content-Type": "application/scim+json"} SCIM_CORE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:User" SCIM_CORE_GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group" diff --git a/tests/integration/management/test_scim_group_bulk_roster.py b/tests/integration/management/test_scim_group_bulk_roster.py index 2af6ae0a83d..6d8e657c5ce 100644 --- a/tests/integration/management/test_scim_group_bulk_roster.py +++ b/tests/integration/management/test_scim_group_bulk_roster.py @@ -36,6 +36,8 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from psycopg import sql from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group" PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp" SMALL_PUSH: Final = 5 diff --git a/tests/integration/management/test_scim_group_pathless_patch.py b/tests/integration/management/test_scim_group_pathless_patch.py index efd3afba5b3..d3bdfb840c3 100644 --- a/tests/integration/management/test_scim_group_pathless_patch.py +++ b/tests/integration/management/test_scim_group_pathless_patch.py @@ -15,6 +15,8 @@ from integration._support.process import owned_proxy_process from integration.authorization._guardrail_opt_out import upstream_hits from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + PATCH_OP_SCHEMA: Final = "urn:ietf:params:scim:api:messages:2.0:PatchOp" GROUP_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:Group" BURST_TEAMS: Final = 10 diff --git a/tests/integration/management/test_search_tool_crud.py b/tests/integration/management/test_search_tool_crud.py index d7f161431c6..f46da964e14 100644 --- a/tests/integration/management/test_search_tool_crud.py +++ b/tests/integration/management/test_search_tool_crud.py @@ -5,12 +5,15 @@ from types import MappingProxyType from typing import Final import yaml +import pytest from integration._support.client import Gateway, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _ORIGINAL_KEY: Final = "exa-original-provider-key" _ROTATED_KEY: Final = "exa-rotated-provider-key" diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index 628f1807d55..c0c03571a13 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -18,7 +18,7 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, wire_server -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] CLIENT_ID: Final = "integration-sso-client" CLIENT_SECRET: Final = "integration-sso-secret" diff --git a/tests/integration/management/test_team_budget_duration_defaults.py b/tests/integration/management/test_team_budget_duration_defaults.py index fd459f6a7ef..4c3eb71603a 100644 --- a/tests/integration/management/test_team_budget_duration_defaults.py +++ b/tests/integration/management/test_team_budget_duration_defaults.py @@ -10,6 +10,8 @@ from tests.integration._support.client import Gateway, string_value from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy +pytestmark: Final = pytest.mark.shares_worker_state + def _budget_row(team_id: str) -> dict[str, JsonValue]: rows: Final = read_rows( diff --git a/tests/integration/management/test_team_delete_prometheus.py b/tests/integration/management/test_team_delete_prometheus.py index b93e14a0e8d..a089f8ca70d 100644 --- a/tests/integration/management/test_team_delete_prometheus.py +++ b/tests/integration/management/test_team_delete_prometheus.py @@ -24,6 +24,8 @@ from tests.integration._support.client import Gateway, Scenario, eventually, obj from tests.integration._support.database import read_rows from tests.integration._support.process import graceful_stop_seconds, owned_proxy_process +pytestmark: Final = pytest.mark.shares_worker_state + METRIC: Final = "litellm_team_members_metric" METRICS_ROUTE: Final = "/metrics/" MEMBERS: Final = 3 diff --git a/tests/integration/management/test_tool_policy_user.py b/tests/integration/management/test_tool_policy_user.py index 6c0daf50d37..8bc77523aff 100644 --- a/tests/integration/management/test_tool_policy_user.py +++ b/tests/integration/management/test_tool_policy_user.py @@ -22,6 +22,8 @@ from pydantic import BaseModel, JsonValue from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE +pytestmark: Final = pytest.mark.shares_worker_state + AUDIENCE: Final = "litellm-integration" KEY_ID: Final = "integration-signing-key" CLIENT_CLAIM: Final = "client_id" diff --git a/tests/integration/management/test_vector_store_file_list_managed_ids.py b/tests/integration/management/test_vector_store_file_list_managed_ids.py index d5765827969..16d33ee83af 100644 --- a/tests/integration/management/test_vector_store_file_list_managed_ids.py +++ b/tests/integration/management/test_vector_store_file_list_managed_ids.py @@ -26,6 +26,8 @@ from integration._support.ports import allocate_port from openai import APIStatusError, AsyncOpenAI, OpenAI from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + MANAGED_PREFIX: Final = "litellm_proxy:" CARRIED_PROVIDER_FILE_ID: Final = re.compile(r"(?:^|;)llm_output_file_id,([^;]+)") UPLOAD_FILENAME: Final = re.compile(rb'filename="([^"]+)"') diff --git a/tests/integration/mcp/test_mcp_accounting_guardrails.py b/tests/integration/mcp/test_mcp_accounting_guardrails.py index 72d6c860d2f..aff67146ce4 100644 --- a/tests/integration/mcp/test_mcp_accounting_guardrails.py +++ b/tests/integration/mcp/test_mcp_accounting_guardrails.py @@ -36,6 +36,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + DEFAULT_COST: Final = 0.25 ADD_COST: Final = 0.5 FORBIDDEN: Final = "forbidden-integration-word" diff --git a/tests/integration/mcp/test_mcp_agent_365_guardrail.py b/tests/integration/mcp/test_mcp_agent_365_guardrail.py index ee746bf3dbf..2ee3b5e6b89 100644 --- a/tests/integration/mcp/test_mcp_agent_365_guardrail.py +++ b/tests/integration/mcp/test_mcp_agent_365_guardrail.py @@ -27,6 +27,8 @@ from integration._support.mcp import ( from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + TENANT: Final = "00000000-0000-4000-8000-0000000a3650" REJECTED: Final = "Agent 365 guardrail rejected the tool call" GUARDRAIL_ROWS: Final = ( diff --git a/tests/integration/mcp/test_mcp_akto_logging_only.py b/tests/integration/mcp/test_mcp_akto_logging_only.py index 5b3fd84aa83..9b891fcfba2 100644 --- a/tests/integration/mcp/test_mcp_akto_logging_only.py +++ b/tests/integration/mcp/test_mcp_akto_logging_only.py @@ -20,6 +20,8 @@ from integration._support.mcp import EntryPoint, McpCaller, McpPeer, echo_tool, from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + AKTO_KEY: Final = "synthetic-akto-mcp-log-key" BLOCK_MARK: Final = "SYNTHETIC-AKTO-BLOCK" RESPONSE_CHECK: Final = {"akto_connector": "litellm", "response_guardrails": "true", "ingest_data": "true"} diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index e288cef3699..9f6101423bf 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -41,7 +41,7 @@ from openai.types.responses.tool_param import Mcp from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.xdist_group("mcp_llm_private_gateway") +pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] @pytest.fixture(scope="session") diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py index 9775221c6a6..90435f49b34 100644 --- a/tests/integration/mcp/test_mcp_management.py +++ b/tests/integration/mcp/test_mcp_management.py @@ -27,6 +27,8 @@ from pydantic import BaseModel, JsonValue from litellm.models.user import LiteLLM_UserTable from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +pytestmark: Final = pytest.mark.shares_worker_state + ADD: Final = {"a": 4, "b": 5} diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 563d35bf794..327b053d333 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -43,6 +43,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + ADD: Final = {"a": 2, "b": 3} CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb" ACCEPT: Final = {"Accept": "application/json, text/event-stream"} diff --git a/tests/integration/mcp/test_mcp_protocol_errors.py b/tests/integration/mcp/test_mcp_protocol_errors.py index 257e90e9d0e..4930258ec07 100644 --- a/tests/integration/mcp/test_mcp_protocol_errors.py +++ b/tests/integration/mcp/test_mcp_protocol_errors.py @@ -9,6 +9,8 @@ from integration._support.client import Gateway from integration._support.mcp import McpPeer, call_tool, register_mcp, tool_names from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers("other.mcp.errors.protocol_and_malformed_results_cannot_be_empty_success") def test_jsonrpc_error_and_malformed_tool_result_remain_errors(gateway: Gateway) -> None: diff --git a/tests/integration/mcp/test_mcp_resilience.py b/tests/integration/mcp/test_mcp_resilience.py index d821181962f..c335012d787 100644 --- a/tests/integration/mcp/test_mcp_resilience.py +++ b/tests/integration/mcp/test_mcp_resilience.py @@ -32,6 +32,8 @@ from integration._support.process import owned_proxy_process from integration._support.wire import Reply from pydantic import BaseModel, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _ECHO: Final = "catalog-echo:" _PROBE: Final = "catalog-probe" _DESCRIPTION: Final = "Look up one record" diff --git a/tests/integration/mcp/test_mcp_user_env_vars.py b/tests/integration/mcp/test_mcp_user_env_vars.py index 75bb9ede182..246e6ec85d5 100644 --- a/tests/integration/mcp/test_mcp_user_env_vars.py +++ b/tests/integration/mcp/test_mcp_user_env_vars.py @@ -14,6 +14,8 @@ from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, from integration._support.process import owned_proxy_process from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + TOKEN: Final = "USER_TOKEN" WORKSPACE: Final = "WORKSPACE" METHODS: Final = ("GET", "POST", "DELETE") diff --git a/tests/integration/mcp/test_oauth_configuration.py b/tests/integration/mcp/test_oauth_configuration.py index e3a42d630cb..9fd26461286 100644 --- a/tests/integration/mcp/test_oauth_configuration.py +++ b/tests/integration/mcp/test_oauth_configuration.py @@ -19,6 +19,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _Upstream = Callable[[Request], Reply] CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb" diff --git a/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_multi_tool_call_blocks_wire.py b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_multi_tool_call_blocks_wire.py index 1ee19fae339..b7766db946e 100644 --- a/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_multi_tool_call_blocks_wire.py +++ b/tests/integration/messages_endpoint/chat_bridge/test_chat_bridge_multi_tool_call_blocks_wire.py @@ -26,6 +26,8 @@ from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall from openai.types.responses import FunctionToolParam, ResponseFunctionToolCall, ResponseOutputItemDoneEvent from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "mistral-multi-tool" _PROVIDER_MODEL: Final = f"mistral/{_BACKEND}" _PROVIDER_KEY: Final = "integration-provider-key" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py index b13c34fe2d6..e159573e9fd 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py @@ -14,6 +14,8 @@ from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "claude-sonnet-4-6" _API_KEY: Final = "synthetic-anthropic-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py index 8944d1990d9..9dd1a2d2c18 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_chaos.py @@ -45,7 +45,7 @@ _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _OUTAGE_STATUSES: Final = (529, 429, 503) _ROUTING_ENCODED_ID: Final = re.compile(r"resp_([A-Za-z0-9+/]+=*)") -pytestmark = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] Endpoint = Literal["messages", "chat", "responses"] diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py index cabde4f1e92..cd11644eec0 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -20,6 +20,8 @@ from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + BETA: Final = "dangerous-tool-use-2026-09-03" SAFEGUARDS: Final[list[JsonValue]] = [ {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_passthrough_migration_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_passthrough_migration_wire.py index 834e10e32b8..cff6d5ef9b2 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_passthrough_migration_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_passthrough_migration_wire.py @@ -6,11 +6,15 @@ from hashlib import sha256 from pathlib import Path from typing import Final +import pytest + from integration._support.client import Gateway, JsonValue, eventually from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL: Final = "claude-sonnet-4-5-20250929" _KEY: Final = "synthetic-anthropic-key" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py index 63325e7a005..6e4842966a8 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py @@ -25,6 +25,8 @@ from integration._support.wire import Reply, Request, wire_server from integration.providers._cache_control_marks_support import anthropic_peer, owned_config from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _MODEL: Final = "burst-claude" _PROVIDER_KEY: Final = "synthetic-provider-key" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py index a6f098cf64f..4e5ef521c0e 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py @@ -180,6 +180,8 @@ import httpx import pytest from integration._support.client import Gateway, eventually +pytestmark: Final = pytest.mark.shares_worker_state + _QUERY: Final = "integration capped search" _TEXT_BLOCK: Final = {"type": "text", "text": "searching once more"} _NOT_INTERCEPTED: Final = "native tool reached the provider" diff --git a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py index 15c0593f80d..70131aa1468 100644 --- a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py +++ b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py @@ -15,6 +15,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + CLAUDE_MODEL: Final = "caller-tools-claude" OPENAI_MODEL: Final = "caller-tools-openai" GEMINI_MODEL: Final = "caller-tools-gemini" diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index 61def2a355b..a57e25f88ae 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -5,7 +5,7 @@ import json import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from itertools import chain from pathlib import Path from typing import Final @@ -854,11 +854,7 @@ def _rig( if fresh_client_connections else owned.gateway.client ) - proxy: Final = Gateway( - client=proxy_client, - key=owned.gateway.key, - upstream_url=owned.gateway.upstream_url, - ) + proxy: Final = replace(owned.gateway, client=proxy_client) scenario: Final = resources.enter_context(proxy.scenario()) model: Final = scenario.model( model=model_name, diff --git a/tests/integration/observability/azure_dispatch_support.py b/tests/integration/observability/azure_dispatch_support.py index 2f1d758ead4..c54293e90c8 100644 --- a/tests/integration/observability/azure_dispatch_support.py +++ b/tests/integration/observability/azure_dispatch_support.py @@ -10,7 +10,7 @@ from typing import Final import yaml from integration._support.client import Gateway, object_value -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process from integration._support.redis_process import OwnedRedis, owned_redis from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue @@ -528,6 +528,7 @@ def dispatch_proxy( }, config=config, workers=workers, + state=OwnedState.SHARED_WITH_WORKER, ) as owned: yield owned diff --git a/tests/integration/observability/test_akto_logging_only.py b/tests/integration/observability/test_akto_logging_only.py index 31adacdec75..e6cb5b5c8df 100644 --- a/tests/integration/observability/test_akto_logging_only.py +++ b/tests/integration/observability/test_akto_logging_only.py @@ -26,6 +26,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + LOG_KEY: Final = "synthetic-akto-log-key" INPUT_KEY: Final = "synthetic-akto-input-key" OUTPUT_KEY: Final = "synthetic-akto-output-key" diff --git a/tests/integration/observability/test_azure_content_safety_audit.py b/tests/integration/observability/test_azure_content_safety_audit.py index 85ced5f719b..dc0346aee75 100644 --- a/tests/integration/observability/test_azure_content_safety_audit.py +++ b/tests/integration/observability/test_azure_content_safety_audit.py @@ -15,6 +15,8 @@ from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _ATTACK_MARKER: Final = "synthetic-attack-marker" _MODERATION_MARKER: Final = "synthetic-moderation-marker" diff --git a/tests/integration/observability/test_azure_content_safety_endpoints.py b/tests/integration/observability/test_azure_content_safety_endpoints.py index cd52122267f..dd2efb6ad87 100644 --- a/tests/integration/observability/test_azure_content_safety_endpoints.py +++ b/tests/integration/observability/test_azure_content_safety_endpoints.py @@ -13,6 +13,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _ATTACK_MARKER: Final = "synthetic-attack-marker" _AZURE_TARGET_PREFIX: Final = "/contentsafety/text:shieldPrompt?api-version=" diff --git a/tests/integration/observability/test_azure_storage_client_ttl.py b/tests/integration/observability/test_azure_storage_client_ttl.py index f8f32820daa..1890fad3589 100644 --- a/tests/integration/observability/test_azure_storage_client_ttl.py +++ b/tests/integration/observability/test_azure_storage_client_ttl.py @@ -4,6 +4,8 @@ from collections.abc import Callable from pathlib import Path from typing import Final +import pytest + from _azure_storage_support import ( SINK_HOSTS, RecordingDataLakeSink, @@ -18,6 +20,8 @@ from integration._support.process import owned_proxy from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + WORKERS: Final = 2 FLUSH_SECONDS: Final = "1" diff --git a/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py index c55bd181ecb..9546b9549ff 100644 --- a/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py +++ b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py @@ -36,6 +36,8 @@ from integration._support.process import owned_proxy_process from integration._support.upstream import aws_event_stream_frame from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" ROUTER_MODEL: Final = "converse-relay" EVENT_STREAM: Final = "application/vnd.amazon.eventstream" diff --git a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py index 3ef73f613fe..ae243324e27 100644 --- a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py +++ b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py @@ -38,6 +38,8 @@ from integration._support.process import owned_proxy_process from integration._support.upstream import aws_event_stream_frame from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + CLAUDE_ID: Final = "anthropic.claude-sonnet-5-v1:0" LLAMA_ID: Final = "meta.llama4-maverick-17b-instruct-v1:0" PROFILE_ID: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/relayprofile" diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 2ca9e0c6410..c52649c46df 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -18,6 +18,8 @@ from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from prometheus_client.parser import text_string_to_metric_families +pytestmark: Final = pytest.mark.shares_worker_state + GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" DEPLOYMENT_FAILURE: Final = "litellm_deployment_failure_responses_total" DEPLOYMENT_REQUESTS: Final = "litellm_deployment_total_requests_total" diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py index d4b62b36e8b..b6014c72a2f 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Final import httpx +import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.process import owned_proxy_process from integration._support.redis_process import owned_redis @@ -24,6 +25,8 @@ from test_cache_hit_guardrail_metrics import ( _rig, ) +pytestmark: Final = pytest.mark.shares_worker_state + BURST: Final = 10 diff --git a/tests/integration/observability/test_callback_delivery.py b/tests/integration/observability/test_callback_delivery.py index e786e655a13..f3e8daab014 100644 --- a/tests/integration/observability/test_callback_delivery.py +++ b/tests/integration/observability/test_callback_delivery.py @@ -14,6 +14,8 @@ from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers( "other.observability.callbacks.credentials_stay_out_of_event_bodies", diff --git a/tests/integration/observability/test_content_filter_logging_only.py b/tests/integration/observability/test_content_filter_logging_only.py index ab1779c497d..a21c380c177 100644 --- a/tests/integration/observability/test_content_filter_logging_only.py +++ b/tests/integration/observability/test_content_filter_logging_only.py @@ -27,6 +27,8 @@ from integration._support.otlp_sink import AuditConfigWriter, SpanSinks, recorde from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 FORBIDDEN_WORD: Final = "synthetic-forbidden-word" diff --git a/tests/integration/observability/test_custom_logger_redaction.py b/tests/integration/observability/test_custom_logger_redaction.py index adc0eb9103c..2abaa85df25 100644 --- a/tests/integration/observability/test_custom_logger_redaction.py +++ b/tests/integration/observability/test_custom_logger_redaction.py @@ -28,7 +28,7 @@ INSTRUCTIONS: Final = "Reply in French. redaction-instructions-canary" BRIDGED: Final = "bridged-redaction-canary" USAGE_DEPLOYMENT: Final = f"redaction-usage-{uuid.uuid4().hex}" REDACT: Final = {"x-litellm-enable-message-redaction": "true"} -pytestmark = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] def _provider(request: Request) -> Reply: diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4e0c3d45cdf..b38ad2a3873 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -27,6 +27,8 @@ from openai import APIError, AsyncOpenAI, OpenAI from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers("other.observability.guardrails.rewrite_reaches_correct_anthropic_positions") def test_guardrail_rewrites_system_and_user_in_actual_anthropic_request(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/integration/observability/test_guardrail_stream_scope.py b/tests/integration/observability/test_guardrail_stream_scope.py index 8b832f07aa4..071ae2dc5f9 100644 --- a/tests/integration/observability/test_guardrail_stream_scope.py +++ b/tests/integration/observability/test_guardrail_stream_scope.py @@ -25,6 +25,8 @@ from pydantic import TypeAdapter import litellm +pytestmark: Final = pytest.mark.shares_worker_state + JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) EndpointScope: TypeAlias = Literal["streaming", "non_streaming"] ModelKind: TypeAlias = Literal["router", "direct"] diff --git a/tests/integration/observability/test_guardrail_stream_scope_chaos.py b/tests/integration/observability/test_guardrail_stream_scope_chaos.py index 3ada2fb7753..eb534161ab4 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_chaos.py +++ b/tests/integration/observability/test_guardrail_stream_scope_chaos.py @@ -25,6 +25,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + ChaosEndpoint: TypeAlias = Literal["chat", "messages", "responses"] CHAOS_ENDPOINTS: Final[tuple[ChaosEndpoint, ...]] = ("chat", "messages", "responses") CHAOS_MODELS: Final = MappingProxyType( diff --git a/tests/integration/observability/test_guardrail_stream_scope_matrix.py b/tests/integration/observability/test_guardrail_stream_scope_matrix.py index ec8eae66d10..d100e44889c 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_matrix.py +++ b/tests/integration/observability/test_guardrail_stream_scope_matrix.py @@ -38,6 +38,8 @@ from integration.cost_calculation.cost_tracking_case import RealtimeResponse from pydantic import TypeAdapter from websockets.asyncio.server import ServerConnection, serve +pytestmark: Final = pytest.mark.shares_worker_state + Endpoint: TypeAlias = Literal["chat", "messages", "responses"] Mode: TypeAlias = Literal["pre_call", "during_call", "post_call", "logging_only"] Scope: TypeAlias = Literal["streaming", "non_streaming"] diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 7fc2ced63e8..6c66fb0f864 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -33,6 +33,8 @@ from opentelemetry.proto.common.v1.common_pb2 import KeyValue from opentelemetry.proto.trace.v1.trace_pb2 import Span from pydantic import BaseModel, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + PUBLIC_KEY: Final = "pk-lf-integration" SECRET_KEY: Final = "sk-lf-integration" PROJECTS_PATH: Final = "/api/public/projects" diff --git a/tests/integration/observability/test_langtrace_delivery.py b/tests/integration/observability/test_langtrace_delivery.py index f81780f9f72..4709049ef92 100644 --- a/tests/integration/observability/test_langtrace_delivery.py +++ b/tests/integration/observability/test_langtrace_delivery.py @@ -7,7 +7,7 @@ import uuid from collections.abc import Callable, Iterator, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import repeat from pathlib import Path from queue import SimpleQueue @@ -659,7 +659,7 @@ def test_langtrace_survives_a_killed_worker(gateway: Gateway, tmp_path: Path) -> limits=httpx.Limits(max_keepalive_connections=0), ) as fresh_connections, ): - proxy: Final = Gateway(fresh_connections, owned.gateway.key, owned.gateway.upstream_url) + proxy: Final = replace(owned.gateway, client=fresh_connections) with proxy.scenario() as scenario: rig: Final = _Rig(proxy, scenario.model(api_base=provider.url + "/v1"), provider, _Sink(sink, key)) _assert_kill_and_recovery(owned, rig) diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index 48e9930d948..bc8f7bb9da2 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -49,6 +49,8 @@ from integration._support.ports import allocate_port from openai import APIConnectionError as OpenAIAPIConnectionError from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _EDGE_DELAY_SECONDS: Final = 5 diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index 354e3b464fe..7c976ed25b3 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -48,6 +48,8 @@ from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + def test_G7_database_invalid_scope_reads_preserve_pre_call_blocking(gateway: Gateway, tmp_path: Path) -> None: identity: Final = f"logging-scope-g7-{uuid.uuid4().hex}" diff --git a/tests/integration/observability/test_logging_only_scope_runtime.py b/tests/integration/observability/test_logging_only_scope_runtime.py index b637859da52..83bebb64960 100644 --- a/tests/integration/observability/test_logging_only_scope_runtime.py +++ b/tests/integration/observability/test_logging_only_scope_runtime.py @@ -58,6 +58,8 @@ from pydantic import JsonValue from tests.integration.cost_calculation.cost_tracking_case import JsonResponse +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.parametrize( ("row_id", "endpoint", "stream", "client_kind", "scope", "include_scope", "block_directions"), diff --git a/tests/integration/observability/test_otel_cache_hit_cost.py b/tests/integration/observability/test_otel_cache_hit_cost.py index 76c21b529ef..7f1ba11f566 100644 --- a/tests/integration/observability/test_otel_cache_hit_cost.py +++ b/tests/integration/observability/test_otel_cache_hit_cost.py @@ -19,6 +19,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from redis import Redis +pytestmark: Final = pytest.mark.shares_worker_state + INPUT_PRICE: Final = 0.001 OUTPUT_PRICE: Final = 0.002 PROMPT_TOKENS: Final = 7 diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index 6d0e5ea197d..af893abab48 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -23,6 +23,8 @@ from integration._support.process import OwnedProxy, owned_proxy_process, stop_r from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + MARKER: Final = re.compile(rb"otelconv-[0-9a-f]{32}") CONVERSATION: Final = "gen_ai.conversation.id" diff --git a/tests/integration/observability/test_otel_trace_response_header.py b/tests/integration/observability/test_otel_trace_response_header.py index e84aeab48cf..5cb66fded99 100644 --- a/tests/integration/observability/test_otel_trace_response_header.py +++ b/tests/integration/observability/test_otel_trace_response_header.py @@ -23,6 +23,8 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + MARKER: Final = re.compile(rb"traceresp-[0-9a-f]{32}") TRACERESPONSE: Final = re.compile(r"00-([0-9a-f]{32})-([0-9a-f]{16})-01") SERVER: Final = 2 diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index 68ad3d50a80..cd063c469f1 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -18,6 +18,8 @@ from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} _NOT_FOUND_BODY: Final = json.dumps( { diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py index 83560f863ed..47ccf45f829 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_visibility.py +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -18,6 +18,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, NotFoundError, OpenAI from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _UPSTREAM_ERROR: Final[dict[str, JsonValue]] = { "error": { "code": 404, diff --git a/tests/integration/observability/test_s3_v2_flush_surfaces.py b/tests/integration/observability/test_s3_v2_flush_surfaces.py index 5ee117a4b80..1e373171de6 100644 --- a/tests/integration/observability/test_s3_v2_flush_surfaces.py +++ b/tests/integration/observability/test_s3_v2_flush_surfaces.py @@ -22,6 +22,8 @@ from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy from integration._support.wire import wire_server +pytestmark: Final = pytest.mark.shares_worker_state + PER_REQUEST_KEY: Final = re.compile(rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/.+\.json$") BATCH_KEY: Final = re.compile( rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$" diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index 382463eb46c..8703c1f318c 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -33,6 +33,8 @@ from integration._support.database_relay import database_relay from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + FLUSH: Final = {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"} HOUR: Final = {"s3_partition_granularity": "hour"} ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index 674087eadf3..6e50a9bedf5 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -26,6 +26,8 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTrace from opentelemetry.proto.common.v1.common_pb2 import AnyValue from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + MARKER: Final = re.compile(rb"signoz-[0-9a-f]{32}") JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) RESPONSE_ID: Final = "gen_ai.response.id" diff --git a/tests/integration/observability/test_straiker_v3_platform.py b/tests/integration/observability/test_straiker_v3_platform.py index c078cbb0c61..c74b094ee95 100644 --- a/tests/integration/observability/test_straiker_v3_platform.py +++ b/tests/integration/observability/test_straiker_v3_platform.py @@ -31,6 +31,8 @@ from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_pr from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port +pytestmark: Final = pytest.mark.shares_worker_state + V3_KEY: Final = "sk_agt_synthetic_integration_key" V1_KEY: Final = "synthetic-v1-collection-key" V3_PATH: Final = "/api/v3/detect" diff --git a/tests/integration/observability/test_stream_error_alert_api_base.py b/tests/integration/observability/test_stream_error_alert_api_base.py index 449e378dfc7..d7f325da8c7 100644 --- a/tests/integration/observability/test_stream_error_alert_api_base.py +++ b/tests/integration/observability/test_stream_error_alert_api_base.py @@ -34,7 +34,7 @@ from openai.types.chat import ChatCompletionChunk from openai.types.responses import ResponseFailedEvent, ResponseTextDeltaEvent from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] AZURE_KEY: Final = "synthetic-azure-key" API_VERSION: Final = "2025-04-01-preview" diff --git a/tests/integration/observability/test_wingback_guardrail.py b/tests/integration/observability/test_wingback_guardrail.py index e93d873345c..9206acce81d 100644 --- a/tests/integration/observability/test_wingback_guardrail.py +++ b/tests/integration/observability/test_wingback_guardrail.py @@ -13,6 +13,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + API_KEY: Final = "wbk_eg_synthetic" APP_ID: Final = "synthetic-app" DENIED: Final = "synthetic wingback denied marker" diff --git a/tests/integration/pricing/test_configured_prices.py b/tests/integration/pricing/test_configured_prices.py index b7eaadc05f9..8e97c9c8322 100644 --- a/tests/integration/pricing/test_configured_prices.py +++ b/tests/integration/pricing/test_configured_prices.py @@ -20,6 +20,8 @@ from tests.integration._support.upstream import delete_scenario, register_scenar from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse from tests.integration.pricing.test_realtime_cached_audio_pricing import one_realtime_turn +pytestmark: Final = pytest.mark.shares_worker_state + REALTIME_MODEL: Final = "gpt-realtime-2" REALTIME_INPUT_TEXT_TOKENS: Final = 10 REALTIME_INPUT_AUDIO_TOKENS: Final = 20 diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py index bb24130aa6c..0aa3a3bcba2 100644 --- a/tests/integration/pricing/test_price_data_reload_prices.py +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -18,6 +18,8 @@ from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _ROW: Final = "anthropic.claude-audit-reload-v1:0" _GEO: Final = f"au.{_ROW}" _FINETUNE: Final = "ft:gpt-audit-reload-2025-01-01" diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py index 5fccd9804a6..a10017425f1 100644 --- a/tests/integration/providers/_idle_upstream_support.py +++ b/tests/integration/providers/_idle_upstream_support.py @@ -24,7 +24,7 @@ from integration._support.client import Gateway, eventually, held, list_value, o from integration._support.database import read_rows from integration._support.idle_gateway import MARKER, Arrival, Drop, IdleGateway, idle_gateway from integration._support.openai_wire import chat_reply, chat_stream, responses_reply -from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, WorkerHealthcheck, owned_proxy_process from integration._support.wire import Reply, Request from openai.types.chat import ChatCompletionChunk from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent @@ -704,7 +704,13 @@ def proxy_for( directory, tuple(row.deploy(row.model, gateways[row.name].url, steady) for row in rows) ) with owned_proxy_process( - gateway, directory, overrides, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + gateway, + directory, + overrides, + config=config, + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + state=OwnedState.SHARED_WITH_WORKER, ) as owned: yield UpstreamProxy(owned, two_workers(owned), MappingProxyType({row.name: row for row in rows}), gateways) diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py index 461839ccbcc..b47dd242429 100644 --- a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -61,6 +61,7 @@ from pydantic import BaseModel, JsonValue pytestmark: Final = [ pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.xdist_group("anthropic_refusal_stop_details_wire"), + pytest.mark.shares_worker_state, ] _ANTHROPIC_BASE: Final = "http://api.anthropic.com" diff --git a/tests/integration/providers/test_anthropic_thinking_signature_logging_wire.py b/tests/integration/providers/test_anthropic_thinking_signature_logging_wire.py index 6622bbb1c9a..21971335b42 100644 --- a/tests/integration/providers/test_anthropic_thinking_signature_logging_wire.py +++ b/tests/integration/providers/test_anthropic_thinking_signature_logging_wire.py @@ -37,7 +37,7 @@ from integration._support.process import owned_proxy from integration._support.wire import Wire, wire_server from pydantic import JsonValue -pytestmark = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _ANTHROPIC_KEY: Final = "scripted-anthropic-key" _ANTHROPIC_BASE: Final = "http://api.anthropic.com" diff --git a/tests/integration/providers/test_anthropic_tool_search_result_passthrough_wire.py b/tests/integration/providers/test_anthropic_tool_search_result_passthrough_wire.py index 9bae4fe0052..0ce409fb89d 100644 --- a/tests/integration/providers/test_anthropic_tool_search_result_passthrough_wire.py +++ b/tests/integration/providers/test_anthropic_tool_search_result_passthrough_wire.py @@ -27,7 +27,7 @@ from integration.providers._tool_search_support import ( ) from pydantic import JsonValue -pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _ANTHROPIC_KEY: Final = "scripted-anthropic-key" _ANTHROPIC_BASE: Final = "http://api.anthropic.com" diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py index 2031be5f3be..bb7f685da02 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py @@ -22,6 +22,8 @@ from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _KIMI: Final = "global.moonshotai.kimi-k3" _NOVA: Final = "us.amazon.nova-lite-v1:0" _AWS: Final[dict[str, JsonValue]] = { diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py index b0a19bbe817..3baeebcd74a 100644 --- a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py +++ b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py @@ -34,6 +34,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" _CONVERSE_MODEL: Final = f"bedrock/converse/{_MODEL_ID}" _INVOKE_MODEL: Final = f"bedrock/invoke/{_MODEL_ID}" diff --git a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py index 2e881849dc8..a1432ffd7a0 100644 --- a/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py +++ b/tests/integration/providers/test_bedrock_converse_stream_event_frames_wire.py @@ -28,6 +28,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL_ID: Final = "global.moonshotai.kimi-k3" _CONVERSE_MODEL: Final = f"bedrock/converse/{_MODEL_ID}" _STREAM_TARGET: Final = f"/model/{_MODEL_ID}/converse-stream" diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 7d978bc0561..854da6fce0b 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -31,6 +31,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with +pytestmark: Final = pytest.mark.shares_worker_state + _KIMI_US: Final = "us.moonshotai.kimi-k3" _KIMI_GLOBAL: Final = "global.moonshotai.kimi-k3" _KIMI_BASE: Final = "moonshotai.kimi-k3" diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py index 9c7903497e3..bcbf6d0f6f3 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py @@ -21,6 +21,8 @@ from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "anthropic.claude-haiku-4-5" _API_KEY: Final = "synthetic-mantle-bearer" _CONFIG_MODEL: Final = "bedrock-mantle-claude-chat-chaos" diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 1e599103a27..89ad226761e 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -27,6 +27,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with +pytestmark: Final = pytest.mark.shares_worker_state + _API_KEY: Final = "synthetic-mantle-bearer" _SIGNING_KEY: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") _GPT56: Final = "openai.gpt-5.6-terra" diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py index 02f9e42ff31..1ab6d9fd54f 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py @@ -39,7 +39,7 @@ from integration.providers._cache_control_marks_support import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _SURFACES: Final = ("chat", "chat-stream", "chat-unmarked", "messages", "responses") diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py index 47d28860018..6389d90597e 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py @@ -61,6 +61,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.utils import get_prompt_cache_min_tokens +pytestmark: Final = pytest.mark.shares_worker_state + _OWNED_PROXY_TIMEOUT: Final = 2 * graceful_stop_seconds() + 120 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _LANE_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" diff --git a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py index d2787d49bd0..eed0e3c3641 100644 --- a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py +++ b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py @@ -23,7 +23,7 @@ from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _MODEL: Final = "chatgpt/gpt-5.5" _CONFIG_MODEL: Final = "chatgpt-caller-stream-chaos" diff --git a/tests/integration/providers/test_chatgpt_responses_caller_stream_wire.py b/tests/integration/providers/test_chatgpt_responses_caller_stream_wire.py index d4ae7e0b4ae..57eb30e28fe 100644 --- a/tests/integration/providers/test_chatgpt_responses_caller_stream_wire.py +++ b/tests/integration/providers/test_chatgpt_responses_caller_stream_wire.py @@ -26,7 +26,7 @@ from openai.types.responses import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _MODEL: Final = "chatgpt/gpt-5.5" _NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} diff --git a/tests/integration/providers/test_coralbricks_provider_chaos.py b/tests/integration/providers/test_coralbricks_provider_chaos.py index dc4014fdb75..9ec37da52a9 100644 --- a/tests/integration/providers/test_coralbricks_provider_chaos.py +++ b/tests/integration/providers/test_coralbricks_provider_chaos.py @@ -36,6 +36,8 @@ from integration.providers._coralbricks import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + Endpoint = Literal["chat", "messages", "responses"] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") diff --git a/tests/integration/providers/test_coralbricks_provider_env_wire.py b/tests/integration/providers/test_coralbricks_provider_env_wire.py index 6d8cc766b62..23a68dcd42c 100644 --- a/tests/integration/providers/test_coralbricks_provider_env_wire.py +++ b/tests/integration/providers/test_coralbricks_provider_env_wire.py @@ -27,6 +27,8 @@ from integration.providers._coralbricks import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + ENV_KEY: Final = "cb_env_integration_key" CONFIG_MODEL: Final = "coralbricks-env-config" diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index d3b17044207..d96d6ab1be1 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -15,6 +15,8 @@ from integration._support.upstream import ScenarioHandle, delete_scenario, regis from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL: Final = "ai_decide" _ROUTE: Final = "/api/2.0/ai-functions/ai-decide" _ENDPOINT: Final = "databricks-openjev-qwen35-4b" diff --git a/tests/integration/providers/test_decision_model_guardrail.py b/tests/integration/providers/test_decision_model_guardrail.py index 7f584810df3..67e6f0918d1 100644 --- a/tests/integration/providers/test_decision_model_guardrail.py +++ b/tests/integration/providers/test_decision_model_guardrail.py @@ -23,6 +23,8 @@ from pydantic import JsonValue from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse +pytestmark: Final = pytest.mark.shares_worker_state + _PROBABILITY_ABOVE: Final = 0.99 _PROBABILITY_BELOW: Final = 0.01 _THRESHOLD: Final = 0.5 diff --git a/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py b/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py index d01cbde915e..a7a7c58ceed 100644 --- a/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py +++ b/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py @@ -21,6 +21,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _API_KEY: Final = "synthetic-decisions-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _BOX: Final[dict[str, JsonValue]] = {"top": 0.1, "bottom": 0.4, "left": 0.2, "right": 0.6} diff --git a/tests/integration/providers/test_decisions_chaos.py b/tests/integration/providers/test_decisions_chaos.py index 7bbbb09f638..86edc02bc57 100644 --- a/tests/integration/providers/test_decisions_chaos.py +++ b/tests/integration/providers/test_decisions_chaos.py @@ -21,6 +21,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "pplx-decider-v1-27b" _CONFIG_MODEL: Final = "decisions-chaos" _API_KEY: Final = "synthetic-decisions-key" diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 34ab7f1597f..8aba6841cf9 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -25,7 +25,7 @@ from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(900) +pytestmark: Final = [pytest.mark.timeout(900), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _DROPPED: Final = re.compile(r"original model: (\S+), ignoring and continuing") diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index c88bd0b9273..4f01d79d43b 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -29,7 +29,7 @@ from integration._support.wire import Request from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] Provider: TypeAlias = Literal["chatgpt", "copilot"] Endpoint: TypeAlias = Literal["chat", "responses", "messages"] diff --git a/tests/integration/providers/test_fal_ai_passthrough_wire.py b/tests/integration/providers/test_fal_ai_passthrough_wire.py index b80cc601b3b..b97bd1f2c1c 100644 --- a/tests/integration/providers/test_fal_ai_passthrough_wire.py +++ b/tests/integration/providers/test_fal_ai_passthrough_wire.py @@ -8,6 +8,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL: Final = "fal-ai/trellis-2" _REQUEST_BODY: Final = {"image_url": "https://example.com/in.png", "resolution": 1536} _UPSTREAM_BODY: Final = { diff --git a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py index 17483147107..2be02565f76 100644 --- a/tests/integration/providers/test_gemini_embedding_file_block_chaos.py +++ b/tests/integration/providers/test_gemini_embedding_file_block_chaos.py @@ -19,6 +19,8 @@ from integration._support.process import group_members, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + BACKEND: Final = "gemini-embedding-001" TARGET: Final = f"/models/{BACKEND}:batchEmbedContents" END_OFFSETS: Final = {"fast": "3s", "slow": "5s", "malformed": "3s", "dropped": "9s"} diff --git a/tests/integration/providers/test_gemini_long_decode_chaos.py b/tests/integration/providers/test_gemini_long_decode_chaos.py index 7f86a01d69a..096e30ff814 100644 --- a/tests/integration/providers/test_gemini_long_decode_chaos.py +++ b/tests/integration/providers/test_gemini_long_decode_chaos.py @@ -55,6 +55,8 @@ from integration.providers._gemini_long_decode_support import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + Continuation = Callable[[Request, str], Reply] _MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py index eccfb8e3486..1270b864cb7 100644 --- a/tests/integration/providers/test_gemini_long_decode_wire.py +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -88,6 +88,8 @@ from integration.providers._gemini_long_decode_support import ( ) from pydantic import BaseModel, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + Client = Literal["sdk_sync", "sdk_async", "httpx"] Failure = Literal["400", "500", "drop", "notjson", "errorframe"] Dropped = Literal["first", "continued"] diff --git a/tests/integration/providers/test_gemini_thinking_replay_wire.py b/tests/integration/providers/test_gemini_thinking_replay_wire.py index c0fbe0c4fa3..4f965f079ec 100644 --- a/tests/integration/providers/test_gemini_thinking_replay_wire.py +++ b/tests/integration/providers/test_gemini_thinking_replay_wire.py @@ -26,6 +26,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "gemini-2.5-flash" _API_KEY: Final = "synthetic-gemini-key" _PROJECT: Final = "scripted-project" diff --git a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py index 0194f138689..ead83ed66aa 100644 --- a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py +++ b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py @@ -21,6 +21,8 @@ from integration._support.process import graceful_stop_seconds, owned_proxy_proc from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "gemini-2.5-flash" _API_KEY: Final = "synthetic-gemini-key" _CONFIG_MODEL: Final = "gemini-tool-result-part-reference-chaos" diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py index 3a5a9bf5d06..08ac73cd8e7 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_owned_proxy.py @@ -58,7 +58,7 @@ from integration.providers._gemini_tts_voice_support import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") BURST: Final = 12 diff --git a/tests/integration/providers/test_inferenceapis_context_window_fallbacks_wire.py b/tests/integration/providers/test_inferenceapis_context_window_fallbacks_wire.py index 04869ebb77f..435b4c08c29 100644 --- a/tests/integration/providers/test_inferenceapis_context_window_fallbacks_wire.py +++ b/tests/integration/providers/test_inferenceapis_context_window_fallbacks_wire.py @@ -21,7 +21,7 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter -pytestmark = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)) +pytestmark: Final = [pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)), pytest.mark.shares_worker_state] _BACKEND: Final = "deepseek-ai/DeepSeek-V4-Flash" _VENDOR_KEY: Final = "synthetic-inferenceapis-key" diff --git a/tests/integration/providers/test_malformed_tools_chaos.py b/tests/integration/providers/test_malformed_tools_chaos.py index 0ab50a5afbe..5a8ced9d4bb 100644 --- a/tests/integration/providers/test_malformed_tools_chaos.py +++ b/tests/integration/providers/test_malformed_tools_chaos.py @@ -25,6 +25,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.worker_readiness import wait_for_worker_readiness from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _OPENAI_BACKEND: Final = "gpt-4o-mini" _CLAUDE_BACKEND: Final = "claude-haiku-4-5" _GEMINI_BACKEND: Final = "gemini-2.5-flash" diff --git a/tests/integration/providers/test_ollama_prompt_tools_wire.py b/tests/integration/providers/test_ollama_prompt_tools_wire.py index 59a2ee58a44..428af5dcfe6 100644 --- a/tests/integration/providers/test_ollama_prompt_tools_wire.py +++ b/tests/integration/providers/test_ollama_prompt_tools_wire.py @@ -21,6 +21,8 @@ from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _BACKEND: Final = "llama3-prompt-tools" _API_KEY: Final = "synthetic-ollama-key" _INSTRUCTION: Final = ( diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index 00f2c430ea5..3df73d3b767 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -51,6 +51,8 @@ from integration.providers._mantle_gpt_prompt_cache_support import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _FOUNDRY_BASE: Final = "http://prompt-cache-breakpoint-audit.services.ai.azure.com" _FOUNDRY_HOST: Final = urlsplit(_FOUNDRY_BASE).netloc diff --git a/tests/integration/providers/test_openai_realtime_transcription_url.py b/tests/integration/providers/test_openai_realtime_transcription_url.py index 81c450d8f70..eda8819beb5 100644 --- a/tests/integration/providers/test_openai_realtime_transcription_url.py +++ b/tests/integration/providers/test_openai_realtime_transcription_url.py @@ -36,7 +36,7 @@ from tests.integration.cost_calculation.cost_tracking_case import RealtimeRespon RecordProperty = Callable[[str, object], None] -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = [pytest.mark.timeout(180), pytest.mark.shares_worker_state] TRANSCRIBE_MODEL: Final = "gpt-live-transcribe" SECOND_TRANSCRIBE_MODEL: Final = "gpt-live-transcribe-second" diff --git a/tests/integration/providers/test_openai_responses_tool_search_wire.py b/tests/integration/providers/test_openai_responses_tool_search_wire.py index 0e5c539e3f9..39a955be13b 100644 --- a/tests/integration/providers/test_openai_responses_tool_search_wire.py +++ b/tests/integration/providers/test_openai_responses_tool_search_wire.py @@ -51,7 +51,7 @@ from integration.providers.test_responses_bridge_tool_search_wire import ( ) from pydantic import JsonValue -pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _OPENAI_HOST: Final = "api.openai.com" _OPENAI_API_BASE: Final = f"https://{_OPENAI_HOST}/v1" diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index 4f9a119db2f..267e175fe19 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -22,7 +22,7 @@ from integration._support.responses_ws_peer import ( from integration._support.tls import write_self_signed_cert from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key" PROXY_AZURE_KEY: Final = "proxy-environment-azure-key" diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index c7489685273..c033e390076 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -54,7 +54,7 @@ from tests.integration._support.process import ( from tests.integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] WORKERS: Final = 2 BURST: Final = 20 diff --git a/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py index 3de825ae4a6..d8fb0f8f993 100644 --- a/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py @@ -49,7 +49,7 @@ from integration.providers._cache_control_marks_support import ( ) from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _OWNED: Final = "bridge-mark-owned" diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py index 5acf3b10e04..6fa310c7cf5 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py @@ -30,6 +30,8 @@ from integration.providers._responses_bridge_prompt_cache_breakpoint import ( _send_call, ) +pytestmark: Final = pytest.mark.shares_worker_state + _CONFIG_MODEL: Final = "responses-bridge-cache-breakpoint-chaos" _API_KEY: Final = "synthetic-responses-bridge-key" diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py index 1bc7853437b..f49d7c59710 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py @@ -21,7 +21,7 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _GLOBAL_UNSET: Final = "bridge-breakpoint-global-unset" _GLOBAL_FALSE: Final = "bridge-breakpoint-global-false" diff --git a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py index 6beccefe58f..d4a5e6bf0db 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py @@ -30,6 +30,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + _CLAUDE: Final = cc.OPUS _BEDROCK_CLAUDE: Final = "anthropic.claude-sonnet-5-5" _MISTRAL: Final = "mistral-large-latest" diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py index 75c5efc3aa0..c13d4da7dbb 100644 --- a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -61,7 +61,7 @@ from integration.providers._anthropic_safeguards_beta_support import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.xdist_group("anthropic_safeguards_beta") +pytestmark: Final = [pytest.mark.xdist_group("anthropic_safeguards_beta"), pytest.mark.shares_worker_state] _CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) _BURST: Final = 24 diff --git a/tests/integration/providers/test_sagemaker_embedding_chaos.py b/tests/integration/providers/test_sagemaker_embedding_chaos.py index 8e984c81ff9..ca04aace346 100644 --- a/tests/integration/providers/test_sagemaker_embedding_chaos.py +++ b/tests/integration/providers/test_sagemaker_embedding_chaos.py @@ -31,7 +31,7 @@ from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue -pytestmark = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] _MODEL: Final = "sm-openai-chaos" _ALIASES: Final = ( diff --git a/tests/integration/providers/test_sagemaker_embedding_wire.py b/tests/integration/providers/test_sagemaker_embedding_wire.py index 0ed38b1a768..74f3c8c8a00 100644 --- a/tests/integration/providers/test_sagemaker_embedding_wire.py +++ b/tests/integration/providers/test_sagemaker_embedding_wire.py @@ -37,7 +37,7 @@ from integration._support.sagemaker_embedding import ( from integration._support.wire import Request, Wire, wire_server from pydantic import JsonValue -pytestmark = pytest.mark.timeout(240) +pytestmark: Final = [pytest.mark.timeout(240), pytest.mark.shares_worker_state] _NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} _SPEND_ROW: Final = 'SELECT status, call_type, model_group FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s' diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py index 9b7a9dbbfd0..98ddcca2bae 100644 --- a/tests/integration/providers/test_strands_decider_agentcore_wire.py +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -166,7 +166,7 @@ _SCRUBBED_NAMES: Final = frozenset( _load_yaml: Final[Callable[[str], object]] = yaml.safe_load _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120 -pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +pytestmark: Final = [pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS), pytest.mark.shares_worker_state] _LINEAR_FACTOR: Final = 20 _CLASSIFIER_SAMPLES: Final = 5 diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index bba79049283..cb295815ae2 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -31,6 +31,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _JSON_TOOL: Final = "json_tool_call" _USER_TOOL: Final = "lookup" _STRUCTURED_OUTPUTS_BETA: Final = "structured-outputs-2025-11-13" diff --git a/tests/integration/providers/test_vertex_passthrough_wire.py b/tests/integration/providers/test_vertex_passthrough_wire.py index 06c7e8e03ae..98eb0547d42 100644 --- a/tests/integration/providers/test_vertex_passthrough_wire.py +++ b/tests/integration/providers/test_vertex_passthrough_wire.py @@ -41,7 +41,7 @@ from integration.providers._tool_search_support import ( ) from pydantic import JsonValue, TypeAdapter -pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _PROJECT: Final = "scripted-project" _TOKEN: Final = "scripted-vertex-token" diff --git a/tests/integration/providers/test_voyage_mongodb_host_wire.py b/tests/integration/providers/test_voyage_mongodb_host_wire.py index 62860c52568..4a1f79a609a 100644 --- a/tests/integration/providers/test_voyage_mongodb_host_wire.py +++ b/tests/integration/providers/test_voyage_mongodb_host_wire.py @@ -62,7 +62,7 @@ _JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) _WORKER_PIDS: Final = TypeAdapter(list[int]) _load_yaml: Final[Callable[[str], object]] = yaml.safe_load _OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120 -pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +pytestmark: Final = [pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS), pytest.mark.shares_worker_state] def _deployment(name: str, model: str, mode: str, **litellm_params: JsonValue) -> dict[str, JsonValue]: diff --git a/tests/integration/routing/test_complexity_router_databricks_classifier.py b/tests/integration/routing/test_complexity_router_databricks_classifier.py index af4b76e888f..73a665cee7b 100644 --- a/tests/integration/routing/test_complexity_router_databricks_classifier.py +++ b/tests/integration/routing/test_complexity_router_databricks_classifier.py @@ -30,6 +30,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS +pytestmark: Final = pytest.mark.shares_worker_state + _MODEL: Final = "ai_decide" _ROUTE: Final = "/api/2.0/ai-functions/ai-decide" _API_KEY: Final = "synthetic-databricks-key" diff --git a/tests/integration/routing/test_deployment_mode_routing_wire.py b/tests/integration/routing/test_deployment_mode_routing_wire.py index 65e64883b18..fd7091d9e05 100644 --- a/tests/integration/routing/test_deployment_mode_routing_wire.py +++ b/tests/integration/routing/test_deployment_mode_routing_wire.py @@ -26,6 +26,8 @@ from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _API_KEY: Final = "synthetic-routing-bearer" _MANTLE_GPT: Final = "bedrock_mantle/openai.gpt-5.6-terra" _MANTLE_OSS: Final = "bedrock_mantle/openai.gpt-oss-120b" diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index f30c13993d4..e6ec3444369 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -50,6 +50,8 @@ from integration._support.ports import allocate_port from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse, StoredResponse from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + ORDER_ONE_BLOB: Final = "gAAAAA-minted-by-order-1" ORDER_TWO_BLOB: Final = "gAAAAA-minted-by-order-2" SUMMARY_TEXT: Final = "multiply 17 by 23" diff --git a/tests/integration/routing/test_observed_routing.py b/tests/integration/routing/test_observed_routing.py index d9ac0652a0d..673aa54fb78 100644 --- a/tests/integration/routing/test_observed_routing.py +++ b/tests/integration/routing/test_observed_routing.py @@ -13,6 +13,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers( "other.routing.retries.several_attempts_reach_success_without_hidden_retries", diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index 12190efee26..bc5cf1a5913 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -43,7 +43,7 @@ from integration.routing.test_payment_required_mapping import ( from pydantic import JsonValue from redis import Redis -pytestmark: Final = pytest.mark.timeout(300) +pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] COOLDOWN_SECONDS: Final = 300 RELOAD_SECONDS: Final = 2 diff --git a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py index c9a4e8e8881..a5e93258087 100644 --- a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py +++ b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py @@ -29,6 +29,8 @@ from integration.providers._cache_control_marks_support import anthropic_peer, m from pydantic import JsonValue, TypeAdapter from redis import Redis +pytestmark: Final = pytest.mark.shares_worker_state + _WINDOW: Final = "window-claude" _ITPM: Final = "itpm-claude" _BUDGET: Final = "budget-claude" diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index 1ca92505191..0a758d51b7e 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -59,7 +59,7 @@ from integration.routing._prompt_cache_pin_support import ( from pydantic import JsonValue from redis import Redis -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 300), pytest.mark.shares_worker_state] _CLAUDE: Final = "pin-claude" _SLOW: Final = "pin-slow" diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py b/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py index 4cdda116d23..af418ee9eae 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl_chaos.py @@ -36,7 +36,7 @@ from integration.routing._prompt_cache_pin_support import ( from pydantic import JsonValue from redis import Redis -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 600) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 600), pytest.mark.shares_worker_state] _GROUP: Final = "pin-chaos" _LEFT: Final = "pin-chaos-left" diff --git a/tests/integration/routing/test_ptu_shares.py b/tests/integration/routing/test_ptu_shares.py index 9855394c83c..9f57d31a3a0 100644 --- a/tests/integration/routing/test_ptu_shares.py +++ b/tests/integration/routing/test_ptu_shares.py @@ -45,7 +45,7 @@ from redis import Redis if TYPE_CHECKING: from litellm.llms.azure.ptu_capacity import PTUCapacity -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _SHARED: Final = "ptu-shared" _MIXED: Final = "ptu-mixed" diff --git a/tests/integration/routing/test_ptu_shares_regex_timing.py b/tests/integration/routing/test_ptu_shares_regex_timing.py index 19d5261dcb5..66c1fbad7cb 100644 --- a/tests/integration/routing/test_ptu_shares_regex_timing.py +++ b/tests/integration/routing/test_ptu_shares_regex_timing.py @@ -18,7 +18,7 @@ from integration._support.timing import assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 240), pytest.mark.shares_worker_state] _WORST_GROUP: Final = "ptu-regex-worst" _CONTROL_GROUP: Final = "ptu-regex-unshared-worst" diff --git a/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py index 02e70a82c1c..25ccb6202ed 100644 --- a/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py +++ b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py @@ -23,7 +23,7 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct" diff --git a/tests/integration/security/_sinks.py b/tests/integration/security/_sinks.py index 906384e8760..79aa61dfc90 100644 --- a/tests/integration/security/_sinks.py +++ b/tests/integration/security/_sinks.py @@ -62,7 +62,7 @@ from typing import Final import yaml from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment from integration._support.database import read_rows -from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration.security._canary import Canary, canary @@ -240,7 +240,9 @@ def canary_rig( "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {token}", **(environment or {}), } - with owned_proxy_process(gateway, root, overrides, config=config) as owned: + with owned_proxy_process( + gateway, root, overrides, config=config, state=OwnedState.SHARED_WITH_WORKER + ) as owned: admin = Gateway( owned.gateway.client, overrides.get("LITELLM_MASTER_KEY", owned.gateway.key), owned.gateway.upstream_url ) diff --git a/tests/integration/security/test_fine_tuning_managed_ids.py b/tests/integration/security/test_fine_tuning_managed_ids.py index 59daa8cb820..8f179e5a25f 100644 --- a/tests/integration/security/test_fine_tuning_managed_ids.py +++ b/tests/integration/security/test_fine_tuning_managed_ids.py @@ -19,6 +19,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + MANAGED_MODEL: Final = "gpt-4.1-managed-ids" DEPLOYMENT_KEY: Final = "sk-fine-tuning-managed-deployment" TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n' diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 4408bd147f9..ca18ba8a96a 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -42,7 +42,7 @@ from pydantic import JsonValue from litellm.proxy.spend_tracking.budget_reservation import DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK -pytestmark: Final = pytest.mark.timeout(900) +pytestmark: Final = [pytest.mark.timeout(900), pytest.mark.shares_worker_state] _MODEL: Final = "gemini/gemini-3.8-flash" _INPUT_TOKENS: Final = 300 diff --git a/tests/integration/spend/test_daily_activity_routes.py b/tests/integration/spend/test_daily_activity_routes.py index 73dc95dc47b..520c8ccc6d8 100644 --- a/tests/integration/spend/test_daily_activity_routes.py +++ b/tests/integration/spend/test_daily_activity_routes.py @@ -37,6 +37,8 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) +pytestmark: Final = pytest.mark.shares_worker_state + def _delete_organization(gateway: Gateway, organization_id: str) -> None: response: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]}) diff --git a/tests/integration/spend/test_daily_tag_spend_retention.py b/tests/integration/spend/test_daily_tag_spend_retention.py index 0bdb981c878..2ccd11332a1 100644 --- a/tests/integration/spend/test_daily_tag_spend_retention.py +++ b/tests/integration/spend/test_daily_tag_spend_retention.py @@ -14,6 +14,8 @@ from tests.integration._support.client import Gateway, eventually, string_value from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy, owned_proxy_process +pytestmark: Final = pytest.mark.shares_worker_state + CLEANUP_EVERY_MINUTE: Final = "* * * * *" RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period" _MAPPING: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/spend/test_disable_end_user_cost_tracking.py b/tests/integration/spend/test_disable_end_user_cost_tracking.py index d0233378d1e..0db72a06703 100644 --- a/tests/integration/spend/test_disable_end_user_cost_tracking.py +++ b/tests/integration/spend/test_disable_end_user_cost_tracking.py @@ -22,7 +22,7 @@ from integration._support.upstream import JsonResponse, SseResponse, delete_scen from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] Surface = Literal["chat", "chat-stream", "messages", "messages-stream", "responses", "responses-stream"] diff --git a/tests/integration/spend/test_end_user_spend_without_proxy_user.py b/tests/integration/spend/test_end_user_spend_without_proxy_user.py index 16210438e82..68beedd087d 100644 --- a/tests/integration/spend/test_end_user_spend_without_proxy_user.py +++ b/tests/integration/spend/test_end_user_spend_without_proxy_user.py @@ -13,6 +13,8 @@ from integration._support.process import owned_proxy from integration._support.upstream import JsonResponse, SseResponse, delete_scenario, register_scenario from pydantic import BaseModel, JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/spend/test_expired_ui_session_failure_rows.py b/tests/integration/spend/test_expired_ui_session_failure_rows.py index 0acdc86709e..7bd41eff8e9 100644 --- a/tests/integration/spend/test_expired_ui_session_failure_rows.py +++ b/tests/integration/spend/test_expired_ui_session_failure_rows.py @@ -11,7 +11,7 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from litellm.proxy._types import hash_token -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 240), pytest.mark.shares_worker_state] SESSION_SECONDS: Final = 3 DASHBOARD_INFO_ROUTES: Final = ( diff --git a/tests/integration/spend/test_key_budget_lockout.py b/tests/integration/spend/test_key_budget_lockout.py index 56b4db89bcb..74470b80543 100644 --- a/tests/integration/spend/test_key_budget_lockout.py +++ b/tests/integration/spend/test_key_budget_lockout.py @@ -15,6 +15,8 @@ from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, JsonValue, TypeAdapter from redis import Redis +pytestmark: Final = pytest.mark.shares_worker_state + _CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue]) _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _EXPECTED_MODEL_DISCOVERY: Final = ("GET", "/v1/models") diff --git a/tests/integration/spend/test_messages_stream_usage_cost.py b/tests/integration/spend/test_messages_stream_usage_cost.py index 692dd01e0b0..e28ab121de4 100644 --- a/tests/integration/spend/test_messages_stream_usage_cost.py +++ b/tests/integration/spend/test_messages_stream_usage_cost.py @@ -12,6 +12,8 @@ from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + BEDROCK_MODEL: Final = "anthropic.claude-haiku-4-5-20251001-v1:0" INPUT_TOKENS: Final = 30 CACHE_READ_TOKENS: Final = 900 diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index 407b40d36c3..a364327290f 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -22,6 +22,8 @@ from integration._support.timing import assert_not_delayed_by, assert_within_fac from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + CRAFTED_MODEL: Final = ("exceeded " * 32_000)[:288_000] HOSTILE_5KB_MODEL: Final = ("exceeded budget " * 400)[:5_000] BENIGN_288KB_MODEL: Final = "x" * len(CRAFTED_MODEL) diff --git a/tests/integration/spend/test_passthrough_budget_reservation.py b/tests/integration/spend/test_passthrough_budget_reservation.py index f072333c616..4940e6f61c6 100644 --- a/tests/integration/spend/test_passthrough_budget_reservation.py +++ b/tests/integration/spend/test_passthrough_budget_reservation.py @@ -12,6 +12,8 @@ from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + INPUT_COST_PER_TOKEN: Final = 0.000001 OUTPUT_COST_PER_TOKEN: Final = 0.001 PROMPT_TOKENS: Final = 10 diff --git a/tests/integration/spend/test_passthrough_request_tags.py b/tests/integration/spend/test_passthrough_request_tags.py index f3c0b66b20d..980cb90ff79 100644 --- a/tests/integration/spend/test_passthrough_request_tags.py +++ b/tests/integration/spend/test_passthrough_request_tags.py @@ -15,6 +15,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + def _chat_reply(marker: str) -> dict[str, JsonValue]: return { diff --git a/tests/integration/spend/test_ptu_share_attribution.py b/tests/integration/spend/test_ptu_share_attribution.py index 2b2405c27de..4a2c7fd2417 100644 --- a/tests/integration/spend/test_ptu_share_attribution.py +++ b/tests/integration/spend/test_ptu_share_attribution.py @@ -32,7 +32,7 @@ from litellm.constants import PTU_SENTINEL_API_KEY if TYPE_CHECKING: from litellm.llms.azure.ptu_capacity import PTUCapacity -pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 240) +pytestmark: Final = [pytest.mark.timeout(4 * graceful_stop_seconds() + 240), pytest.mark.shares_worker_state] _GROUP: Final = "ptu-ledger" _SHARED_DEPLOYMENT: Final = "ledger-shared-gpt41" diff --git a/tests/integration/spend/test_requester_ip_address_logging.py b/tests/integration/spend/test_requester_ip_address_logging.py index 23e3a56db10..74af0afb59c 100644 --- a/tests/integration/spend/test_requester_ip_address_logging.py +++ b/tests/integration/spend/test_requester_ip_address_logging.py @@ -24,6 +24,8 @@ from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + FORWARDED_IP: Final = "203.0.113.74" PROXY_HOP: Final = "198.51.100.23" REAL_IP: Final = "203.0.113.75" diff --git a/tests/integration/spend/test_reset_budget_leader_election.py b/tests/integration/spend/test_reset_budget_leader_election.py index f9c165a0be0..58c018c4181 100644 --- a/tests/integration/spend/test_reset_budget_leader_election.py +++ b/tests/integration/spend/test_reset_budget_leader_election.py @@ -11,6 +11,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from redis import Redis +pytestmark: Final = pytest.mark.shares_worker_state + RESET_LEASE_KEY: Final = "cronjob_lock:reset_budget_job" PEER_POD_LEASE: Final = json.dumps("integration-peer-pod-holding-the-reset-lease") FAST_RESET_TICK: Final = MappingProxyType( diff --git a/tests/integration/spend/test_responses_websocket_turn_rows.py b/tests/integration/spend/test_responses_websocket_turn_rows.py index 59b5925c6eb..4d72c5b17cf 100644 --- a/tests/integration/spend/test_responses_websocket_turn_rows.py +++ b/tests/integration/spend/test_responses_websocket_turn_rows.py @@ -22,7 +22,7 @@ from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, comple from integration._support.tls import write_self_signed_cert from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 180), pytest.mark.shares_worker_state] INPUT_RATE: Final = 0.001 OUTPUT_RATE: Final = 0.002 diff --git a/tests/integration/spend/test_spend_log_caller_session.py b/tests/integration/spend/test_spend_log_caller_session.py index 213583d21a1..756e02a81ee 100644 --- a/tests/integration/spend/test_spend_log_caller_session.py +++ b/tests/integration/spend/test_spend_log_caller_session.py @@ -38,7 +38,10 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240) +pytestmark: Final = [ + pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240), + pytest.mark.shares_worker_state, +] MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}") CALLER_SESSION_KEY: Final = "litellm_caller_session_id" diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 672c0a38407..9055f6f06c4 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -22,6 +22,8 @@ from pydantic import JsonValue, TypeAdapter from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE from litellm.responses.utils import ResponsesAPIRequestUtils +pytestmark: Final = pytest.mark.shares_worker_state + JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) REDACTED: Final = "REDACTED_BY_LITELM" TOOL_INPUT: Final = {"key": "order-123", "sort_key": "created_at"} diff --git a/tests/integration/spend/test_spend_log_write_batching.py b/tests/integration/spend/test_spend_log_write_batching.py index 386f99b1d07..0db4942ceb1 100644 --- a/tests/integration/spend/test_spend_log_write_batching.py +++ b/tests/integration/spend/test_spend_log_write_batching.py @@ -9,6 +9,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + WRITE_STATEMENT_MAX_BYTES: Final = 200_000 MESSAGES_PER_REQUEST: Final = 60 MESSAGE_CHARACTERS: Final = 2_000 diff --git a/tests/integration/spend/test_spend_logs_metadata_fields.py b/tests/integration/spend/test_spend_logs_metadata_fields.py index 9941edf0f26..cf5b123d1d5 100644 --- a/tests/integration/spend/test_spend_logs_metadata_fields.py +++ b/tests/integration/spend/test_spend_logs_metadata_fields.py @@ -18,6 +18,8 @@ from integration.spend.test_batch_completion_accounting import batch_input_file, from integration.streaming.test_stream_contracts import text_stream from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + CACHED_PROMPT_TOKENS: Final = 4 diff --git a/tests/integration/spend/test_team_member_budget_alerts.py b/tests/integration/spend/test_team_member_budget_alerts.py index f12bcb9748a..be4bbb17e63 100644 --- a/tests/integration/spend/test_team_member_budget_alerts.py +++ b/tests/integration/spend/test_team_member_budget_alerts.py @@ -9,6 +9,8 @@ from integration._support.database import read_rows from integration._support.mail import smtp_sink from integration._support.process import owned_proxy +pytestmark: Final = pytest.mark.shares_worker_state + MEMBER_BUDGET: Final = 0.10 CALL_COST: Final = 20 * 0.001 + 20 * 0.002 diff --git a/tests/integration/spend/test_team_member_spend_flush.py b/tests/integration/spend/test_team_member_spend_flush.py index 2433731f733..e4b897e087e 100644 --- a/tests/integration/spend/test_team_member_spend_flush.py +++ b/tests/integration/spend/test_team_member_spend_flush.py @@ -19,6 +19,8 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + SINGLE_CONNECTION_CONFIG: Final = """ model_list: [] general_settings: diff --git a/tests/integration/spend/test_user_budget_on_team_keys.py b/tests/integration/spend/test_user_budget_on_team_keys.py index f8e0eea3454..f80859555f0 100644 --- a/tests/integration/spend/test_user_budget_on_team_keys.py +++ b/tests/integration/spend/test_user_budget_on_team_keys.py @@ -10,6 +10,8 @@ from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.process import owned_proxy +pytestmark: Final = pytest.mark.shares_worker_state + @pytest.mark.covers("spend.user_budget.opted_in_team_key_is_denied_once_owner_budget_is_exhausted") def test_team_key_is_denied_before_provider_once_owner_personal_budget_is_exhausted_when_opted_in( diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index 9aedb81c3e0..6900a6d489a 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -28,6 +28,8 @@ from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _IMAGE_BACKEND: Final = "gemini-scripted-image" _CHAT_BACKEND: Final = "gemini-scripted-chat" _PROJECT: Final = "scripted-project" diff --git a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py index 021ca68ce8e..a8ebe510b8e 100644 --- a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py +++ b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py @@ -23,6 +23,7 @@ from litellm.proxy import client as proxy_client pytestmark: Final = [ pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.xdist_group("chat-stream-guardrail-error-events"), + pytest.mark.shares_worker_state, ] _MODEL: Final = "stream-error-events" diff --git a/tests/integration/streaming/test_faked_responses_stream_chaos.py b/tests/integration/streaming/test_faked_responses_stream_chaos.py index fca99c9d525..612451c237e 100644 --- a/tests/integration/streaming/test_faked_responses_stream_chaos.py +++ b/tests/integration/streaming/test_faked_responses_stream_chaos.py @@ -41,7 +41,7 @@ from integration.streaming.faked_responses_case import ( ) from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _RESPONSES_MODEL: Final = "faked-stream-chaos" _BRIDGE_MODEL: Final = "faked-stream-chaos-bridge" diff --git a/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py b/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py index 796e706c569..1cb778723fd 100644 --- a/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py +++ b/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py @@ -18,6 +18,8 @@ from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _PROJECT: Final = "scripted-project" _ADAPTER_MODEL: Final = "settings-adapter-openai" diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 83a3af06355..0ca62d2f014 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -36,7 +36,7 @@ from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEven from pydantic import JsonValue _CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 -pytestmark: Final = pytest.mark.timeout(_CELL_SECONDS) +pytestmark: Final = [pytest.mark.timeout(_CELL_SECONDS), pytest.mark.shares_worker_state] _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) diff --git a/tests/integration/streaming/test_noma_v2_streaming_settings.py b/tests/integration/streaming/test_noma_v2_streaming_settings.py index 6354a2c9135..877978732f8 100644 --- a/tests/integration/streaming/test_noma_v2_streaming_settings.py +++ b/tests/integration/streaming/test_noma_v2_streaming_settings.py @@ -21,6 +21,8 @@ from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_pr from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue +pytestmark: Final = pytest.mark.shares_worker_state + Endpoint: TypeAlias = Literal["chat", "messages", "responses"] Spelling: TypeAlias = Literal["noma_v2", "noma_use_v2", "stored", "custom"] Variant: TypeAlias = Literal[ diff --git a/tests/integration/streaming/test_responses_bridge_stream_chaos.py b/tests/integration/streaming/test_responses_bridge_stream_chaos.py index 754ee5d3981..839a19b64a8 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_chaos.py +++ b/tests/integration/streaming/test_responses_bridge_stream_chaos.py @@ -37,7 +37,7 @@ from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] _MODEL: Final = "bridged-stream-chaos" _MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index a3efbd233b9..8a5655be0f8 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -50,7 +50,7 @@ from pydantic import JsonValue _PRIMARY: Final = "bridged-primary" _SPARE: Final = "bridged-spare" -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] def assert_failed_once(wire: Wire, call_id: str, model: str, attempts: int = 1) -> tuple[Request, ...]: diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index 3f8243033d7..d7537a78db8 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -18,6 +18,8 @@ from integration._support.wire import Reply, Request, wire_server from openai import OpenAI from pydantic import JsonValue, TypeAdapter +pytestmark: Final = pytest.mark.shares_worker_state + JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) diff --git a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py index 194e6972453..efd649f9383 100644 --- a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py +++ b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py @@ -41,7 +41,7 @@ from pydantic import JsonValue __all__ = ("thoughts_off",) -pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +pytestmark = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT: Final = "**My Inspection Strategy: Validation First**\n\nOkay, here's how I'm approaching this project review. The instruction specifically mentions considering zero, negative, and non-integer quantities. To efficiently address that, I'm going to start with `shop/validation.py`. My reasoning is simple: the validation logic is the first line of defense. Understanding how the system *expects* to handle these edge cases will inform everything else. It's the logical place to begin my investigation to understand the code. I'll use the `read_file` tool to examine that file. I am **not** changing any files. I will call the `chat_completions` client for this project review. I am now proceeding with reading `shop/validation.py`.\n\n\n" diff --git a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py index cb6f973d2df..1439ee3fed1 100644 --- a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py +++ b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py @@ -21,6 +21,8 @@ from integration.translation.responses.bases.azure_ai import ( from integration.translation.responses.bases.openai_responses import GPT_5_6_SOL_TEST_CASE, GPT_6_1_SOL_TEST_CASE from integration.translation.runner import assert_translation +pytestmark: Final = pytest.mark.shares_worker_state + GPT_6_1_SOL_TWO_MESSAGES_BEFORE_TOOLS_TEST_CASE: Final = two_messages_before_tools(GPT_6_1_SOL_TEST_CASE) AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( AZURE_CLAUDE_SONNET_5_5_TEST_CASE, diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 7741de93a80..ea9c9fcc854 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -8,11 +8,11 @@ import subprocess import sys import threading import time -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from contextlib import contextmanager from functools import reduce from pathlib import Path -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import Final from unittest.mock import MagicMock @@ -25,13 +25,20 @@ from pydantic import TypeAdapter from litellm.proxy.types_utils.utils import get_instance_fn from tests.integration._support.client import Gateway, eventually from tests.integration._support.process import ( + OWNED_STATE_VARIABLE, OwnedProxy, + OwnedState, + PrivateState, WorkerHealthcheck, WorkerKill, _is_gone, + _proxy_environment, + boots_private_state, healthcheck_seconds, kill_stray_process, + owned_environment, owned_proxy_process, + private_configuration, replacement_pid, stop_process_group, worker_configuration, @@ -385,6 +392,98 @@ def test_owned_runtime_configuration_preserves_custom_worker_startup_hooks( ) +WORKER_STATE: Final = MappingProxyType( + {"DATABASE_URL": "postgresql://admin@127.0.0.1/worker", "REDIS_HOST": "127.0.0.1", "REDIS_PORT": "6379"} +) +PRIVATE_DATABASE: Final = "postgresql://admin@127.0.0.1/owned_private" + + +@contextmanager +def _private_state(directory: Path, inherited: Mapping[str, str]) -> Iterator[PrivateState]: + yield PrivateState( + PRIVATE_DATABASE, + MappingProxyType({"DATABASE_URL": PRIVATE_DATABASE, "REDIS_HOST": "127.0.0.2", "REDIS_PORT": "6390"}), + ) + + +@pytest.mark.parametrize( + ("state", "overrides", "remove_environment", "marked", "private"), + ( + (None, {}, (), False, True), + (None, {"DATABASE_URL": "postgresql://admin@127.0.0.1/scratch"}, (), False, False), + (None, {}, ("DATABASE_URL",), False, False), + (OwnedState.SHARED_WITH_WORKER, {}, (), False, False), + (None, {}, (), True, False), + (OwnedState.PRIVATE, {}, (), True, True), + ), + ids=("default", "explicit-database", "stripped-database", "shared", "marked-module", "private-in-marked-module"), +) +def test_an_owned_proxy_boots_private_state_unless_it_shares_or_brings_its_own_database( + state: OwnedState | None, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + marked: bool, + private: bool, +) -> None: + marker: Final = {OWNED_STATE_VARIABLE: OwnedState.SHARED_WITH_WORKER.value} if marked else {} + inherited: Final = {**WORKER_STATE, **marker} + assert boots_private_state(state, overrides, remove_environment, inherited) is private + + +def test_a_default_owned_proxy_runs_on_its_own_database_and_redis(tmp_path: Path) -> None: + with ( + httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, + owned_environment(tmp_path, {}, (), None, WORKER_STATE, _private_state) as owned, + ): + environment: Final = _proxy_environment( + Gateway(client, "owned-test-key", "http://127.0.0.1:1"), owned.overrides, (), WORKER_STATE + ) + assert environment["DATABASE_URL"] == PRIVATE_DATABASE != WORKER_STATE["DATABASE_URL"] + assert (environment["REDIS_HOST"], environment["REDIS_PORT"]) == ("127.0.0.2", "6390") + assert owned.database_url == PRIVATE_DATABASE + + +@pytest.mark.parametrize( + ("overrides", "remove_environment", "redis"), + ( + ({}, ("REDIS_HOST", "REDIS_PORT"), (None, None)), + ({"REDIS_HOST": "redis.caller", "REDIS_PORT": "7000"}, (), ("redis.caller", "7000")), + ), + ids=("stripped", "caller-owned"), +) +def test_a_private_owned_proxy_leaves_redis_to_a_caller_that_chose_it( + tmp_path: Path, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + redis: tuple[str | None, str | None], +) -> None: + with ( + httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, + owned_environment(tmp_path, overrides, remove_environment, None, WORKER_STATE, _private_state) as owned, + ): + environment: Final = _proxy_environment( + Gateway(client, "owned-test-key", "http://127.0.0.1:1"), owned.overrides, remove_environment, WORKER_STATE + ) + assert environment["DATABASE_URL"] == PRIVATE_DATABASE + assert (environment.get("REDIS_HOST"), environment.get("REDIS_PORT")) == redis + + +@pytest.mark.parametrize( + ("router_settings", "retries"), + (({"allowed_fails": 3}, 0), ({"num_retries": 2}, 2), (None, 0)), + ids=("inherits-worker-default", "keeps-explicit", "no-router-settings"), +) +def test_a_private_owned_proxy_keeps_the_worker_retry_default_unless_its_config_sets_one( + tmp_path: Path, router_settings: dict[str, int] | None, retries: int +) -> None: + source: Final = tmp_path / "source.yaml" + router: Final = {"router_settings": router_settings} if router_settings else {} + source.write_text(yaml.safe_dump({"model_list": [], **router})) + configured: Final = yaml.safe_load(private_configuration(source, tmp_path / "proxy").read_text()) + assert configured["router_settings"]["num_retries"] == retries + assert configured["model_list"] == [] + + _SUPERVISOR: Final = """import multiprocessing import time