|
| 1 | +import asyncio |
| 2 | +import json |
| 3 | +import threading |
| 4 | +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| 5 | +from pathlib import Path |
| 6 | +from typing import Any, Generator, cast |
| 7 | + |
| 8 | +import pytest |
| 9 | + |
| 10 | +from tests.cli.utils.server import ( |
| 11 | + get_free_port, |
| 12 | + start_cli_server_thread, |
| 13 | + start_job_with_env, |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +class TraceServer(ThreadingHTTPServer): |
| 18 | + posts: list[bytes] |
| 19 | + |
| 20 | + |
| 21 | +class TraceHandler(BaseHTTPRequestHandler): |
| 22 | + def do_POST(self) -> None: |
| 23 | + length = int(self.headers.get("content-length", "0")) |
| 24 | + body = self.rfile.read(length) |
| 25 | + cast(TraceServer, self.server).posts.append(body) |
| 26 | + self.send_response(200) |
| 27 | + self.end_headers() |
| 28 | + self.wfile.write(b"OK") |
| 29 | + |
| 30 | + def log_message(self, format: str, *args: object) -> None: |
| 31 | + return |
| 32 | + |
| 33 | + |
| 34 | +@pytest.fixture |
| 35 | +def trace_stub() -> Generator[tuple[TraceServer, str], None, None]: |
| 36 | + port = get_free_port() |
| 37 | + server = TraceServer(("127.0.0.1", port), TraceHandler) |
| 38 | + server.posts = [] |
| 39 | + thread = threading.Thread(target=server.serve_forever, daemon=True) |
| 40 | + thread.start() |
| 41 | + |
| 42 | + try: |
| 43 | + yield server, f"http://127.0.0.1:{port}" |
| 44 | + finally: |
| 45 | + server.shutdown() |
| 46 | + server.server_close() |
| 47 | + |
| 48 | + |
| 49 | +def write_project(project: Path) -> None: |
| 50 | + (project / "entrypoint.py").write_text( |
| 51 | + """from dataclasses import dataclass |
| 52 | +from uipath.tracing import traced |
| 53 | +
|
| 54 | +
|
| 55 | +@dataclass |
| 56 | +class Input: |
| 57 | + message: str |
| 58 | +
|
| 59 | +
|
| 60 | +@dataclass |
| 61 | +class Output: |
| 62 | + message: str |
| 63 | +
|
| 64 | +
|
| 65 | +@traced(name="actual-agent-span") |
| 66 | +def main(input: Input) -> Output: |
| 67 | + return Output(message=f"ok: {input.message}") |
| 68 | +""", |
| 69 | + encoding="utf-8", |
| 70 | + ) |
| 71 | + (project / "uipath.json").write_text( |
| 72 | + json.dumps({"agents": {"main": "entrypoint.py:main"}}, indent=2), |
| 73 | + encoding="utf-8", |
| 74 | + ) |
| 75 | + |
| 76 | + |
| 77 | +async def run_traced_job( |
| 78 | + port: int, |
| 79 | + project: Path, |
| 80 | + trace_url: str, |
| 81 | + job_key: str, |
| 82 | + command: str, |
| 83 | + extra_args: list[str], |
| 84 | +) -> tuple[dict[str, Any], Path]: |
| 85 | + input_file = project / f"{job_key}-input.json" |
| 86 | + output_file = project / f"{job_key}-output.json" |
| 87 | + input_file.write_text(json.dumps({"message": job_key}), encoding="utf-8") |
| 88 | + |
| 89 | + response = await start_job_with_env( |
| 90 | + port, |
| 91 | + job_key, |
| 92 | + command, |
| 93 | + [ |
| 94 | + "main", |
| 95 | + "--input-file", |
| 96 | + str(input_file), |
| 97 | + "--output-file", |
| 98 | + str(output_file), |
| 99 | + *extra_args, |
| 100 | + ], |
| 101 | + { |
| 102 | + "UIPATH_ACCESS_TOKEN": "fake-token", |
| 103 | + "UIPATH_JOB_KEY": job_key, |
| 104 | + "UIPATH_ORGANIZATION_ID": "org-123", |
| 105 | + "UIPATH_TENANT_ID": "tenant-123", |
| 106 | + "UIPATH_TRACE_BASE_URL": trace_url, |
| 107 | + "UIPATH_TRACING_ENABLED": "true", |
| 108 | + "LOG_LEVEL": "DEBUG", |
| 109 | + }, |
| 110 | + working_directory=str(project), |
| 111 | + ) |
| 112 | + return response, output_file |
| 113 | + |
| 114 | + |
| 115 | +def read_json(path: Path) -> Any: |
| 116 | + return json.loads(path.read_text(encoding="utf-8")) |
| 117 | + |
| 118 | + |
| 119 | +@pytest.mark.parametrize( |
| 120 | + ("command", "extra_args"), |
| 121 | + [ |
| 122 | + ("run", []), |
| 123 | + ("debug", ["--attach", "none"]), |
| 124 | + ], |
| 125 | +) |
| 126 | +def test_server_runs_two_traced_jobs_after_trace_shutdown( |
| 127 | + tmp_path: Path, |
| 128 | + trace_stub: tuple[TraceServer, str], |
| 129 | + command: str, |
| 130 | + extra_args: list[str], |
| 131 | +) -> None: |
| 132 | + port = get_free_port() |
| 133 | + start_cli_server_thread(port) |
| 134 | + |
| 135 | + trace_server, trace_url = trace_stub |
| 136 | + project = tmp_path / "project" |
| 137 | + project.mkdir() |
| 138 | + write_project(project) |
| 139 | + |
| 140 | + job1_response, job1_output = asyncio.run( |
| 141 | + run_traced_job(port, project, trace_url, "job-1", command, extra_args) |
| 142 | + ) |
| 143 | + job2_response, job2_output = asyncio.run( |
| 144 | + run_traced_job(port, project, trace_url, "job-2", command, extra_args) |
| 145 | + ) |
| 146 | + |
| 147 | + assert job1_response["success"] is True, job1_response |
| 148 | + assert job2_response["success"] is True, job2_response |
| 149 | + assert read_json(job1_output) == {"message": "ok: job-1"} |
| 150 | + assert read_json(job2_output) == {"message": "ok: job-2"} |
| 151 | + |
| 152 | + result = read_json(project / "__uipath" / "output.json") |
| 153 | + assert result["status"] == "successful" |
| 154 | + assert "error" not in result |
| 155 | + assert trace_server.posts |
0 commit comments