diff --git a/tests/integration/_support/scripted_client.py b/tests/integration/_support/scripted_client.py deleted file mode 100644 index 9502740b1b5..00000000000 --- a/tests/integration/_support/scripted_client.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Client for registering scenarios with the integration upstream.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from typing import Final - -import httpx -from integration._support.scripted_wires import ( - WIRE_MOUNTS, - Scenario, - ScenarioDeleted, - ScenarioRegistered, - Wire, -) - -CONTROL_URL: Final = os.environ.get("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8190").rstrip("/") - - -@dataclass(frozen=True, slots=True) -class ScenarioHandle: - scenario_id: str - wire: Wire - control_url: str - - def api_base(self) -> str: - return f"{self.control_url}/{self.scenario_id}/{self._mount()}" - - def _mount(self) -> str: - return WIRE_MOUNTS[self.wire] - - -def register_scenario(scenario: Scenario) -> ScenarioHandle: - response: Final = httpx.post( - f"{CONTROL_URL}/__scenarios", - json=scenario.model_dump(mode="json"), - trust_env=False, - timeout=15, - ) - response.raise_for_status() - result: Final = ScenarioRegistered.model_validate_json(response.content) - return ScenarioHandle( - scenario_id=result.scenario_id, - wire=scenario.wire, - control_url=CONTROL_URL, - ) - - -def delete_scenario(handle: ScenarioHandle) -> None: - response: Final = httpx.delete( - f"{CONTROL_URL}/__scenarios/{handle.scenario_id}", - trust_env=False, - timeout=15, - ) - response.raise_for_status() - ScenarioDeleted.model_validate_json(response.content) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index c8e77ad513a..b3e6336dcee 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -4,10 +4,12 @@ import argparse from collections import deque import json from dataclasses import dataclass, field +import os from pathlib import Path from queue import SimpleQueue from typing import Final, cast +import httpx import uvicorn from pydantic import JsonValue, TypeAdapter, ValidationError from starlette.applications import Starlette @@ -16,7 +18,16 @@ from starlette.responses import JSONResponse, Response from starlette.routing import Route from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations -from integration._support.scripted_wires import RenderedResponse, Scenario, ScenarioStore, render +from integration._support.scripted_wires import ( + WIRE_MOUNTS, + RenderedResponse, + Scenario, + ScenarioDeleted, + ScenarioRegistered, + ScenarioStore, + Wire, + render, +) JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) INTERNAL_FIELDS: Final = frozenset( @@ -194,6 +205,48 @@ class Provider: ) +CONTROL_URL: Final = os.environ.get("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8190").rstrip("/") + + +@dataclass(frozen=True, slots=True) +class ScenarioHandle: + scenario_id: str + wire: Wire + control_url: str + + def api_base(self) -> str: + return f"{self.control_url}/{self.scenario_id}/{self._mount()}" + + def _mount(self) -> str: + return WIRE_MOUNTS[self.wire] + + +def register_scenario(scenario: Scenario) -> ScenarioHandle: + response: Final = httpx.post( + f"{CONTROL_URL}/__scenarios", + json=scenario.model_dump(mode="json"), + trust_env=False, + timeout=15, + ) + response.raise_for_status() + result: Final = ScenarioRegistered.model_validate_json(response.content) + return ScenarioHandle( + scenario_id=result.scenario_id, + wire=scenario.wire, + control_url=CONTROL_URL, + ) + + +def delete_scenario(handle: ScenarioHandle) -> None: + response: Final = httpx.delete( + f"{CONTROL_URL}/__scenarios/{handle.scenario_id}", + trust_env=False, + timeout=15, + ) + response.raise_for_status() + ScenarioDeleted.model_validate_json(response.content) + + def main() -> None: parser: Final = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=8190) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 66eb373df33..0cbc837c184 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, ConfigDict from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.scripted_client import delete_scenario, register_scenario +from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_matrix import Case, FrontierModel