litellm/.circleci/scripts/run_migration_tests.py

113 lines
3.8 KiB
Python

from __future__ import annotations
import json
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Final
from xml.etree import ElementTree
SUITES: Final = {
"startup": (("test_startup.py",), 12),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
}
def successful_junit(path: Path, expected: int, exit_code: int) -> bool:
if exit_code != 0 or not path.is_file():
return False
try:
root: Final = ElementTree.parse(path).getroot()
except ElementTree.ParseError:
return False
cases: Final = tuple(root.iter("testcase"))
identities: Final = frozenset((case.get("classname"), case.get("name")) for case in cases)
return len(cases) == len(identities) == expected and all(
not any(case.find(tag) is not None for tag in ("failure", "error", "skipped")) for case in cases
)
def output(*command: str) -> str:
return subprocess.check_output(command, text=True, timeout=90).strip()
def record_image() -> None:
source: Final = output("git", "rev-parse", "HEAD")
image: Final = output("docker", "image", "inspect", "litellm-docker-database:ci", "--format", "{{.Id}}")
revision: Final = output(
"docker",
"image",
"inspect",
"litellm-docker-database:ci",
"--format",
'{{index .Config.Labels "org.opencontainers.image.revision"}}',
)
assert re.fullmatch(r"[0-9a-f]{40}", source), "Invalid source revision"
assert revision == source, "Candidate image revision differs from the tested source"
Path("migration-image.json").write_text(
json.dumps(
{
"source_sha": source,
"image_id": image,
"candidate_image": os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""),
}
)
)
def main() -> int:
suite: Final = os.environ["MIGRATION_TEST_SUITE"]
files, expected = SUITES[suite]
metadata: Final = json.loads(Path("migration-image.json").read_text())
assert metadata["source_sha"] == output("git", "rev-parse", "HEAD"), "Image and test source revisions differ"
assert metadata["image_id"] == output(
"docker", "image", "inspect", os.environ["LITELLM_MIGRATION_TEST_IMAGE"], "--format", "{{.Id}}"
), "Loaded image differs from the build output"
assert metadata["candidate_image"] == os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), "Wrong release candidate"
destination: Final = Path(os.environ["MIGRATION_TEST_OUTPUT"])
junit: Final = destination / "junit" / "results.xml"
junit.parent.mkdir(parents=True, exist_ok=True)
result: Final = subprocess.run(
(
sys.executable,
"-m",
"pytest",
*(f"tests/e2e/migrations/{name}" for name in files),
"-vv",
"--tb=short",
"--durations=10",
f"--junitxml={junit}",
"-o",
"addopts=",
"--reruns=0",
),
check=False,
timeout=1200,
)
passed: Final = successful_junit(junit, expected, result.returncode)
(destination / "verdict.json").write_text(
json.dumps(
{
**metadata,
"suite": suite,
"expected_cases": expected,
"passed": passed,
"pytest_exit_code": result.returncode,
"test_revision": metadata["source_sha"],
"workflow_id": os.environ.get("CIRCLE_WORKFLOW_ID", ""),
"job_number": os.environ.get("CIRCLE_BUILD_NUM", ""),
},
indent=2,
)
)
return 0 if passed else 1
if __name__ == "__main__":
if sys.argv[1:] == ["record-image"]:
record_image()
else:
raise SystemExit(main())