mirror of
https://github.com/codeflash-ai/codeflash-internal.git
synced 2026-05-04 18:25:18 +00:00
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
import enum
|
|
from typing import Self
|
|
|
|
from ninja import Schema
|
|
from pydantic import model_validator
|
|
|
|
from aiservice.models.functions_to_optimize import FunctionToOptimize
|
|
|
|
|
|
class TestingMode(enum.Enum):
|
|
BEHAVIOR = "behavior"
|
|
PERFORMANCE = "performance"
|
|
|
|
|
|
class TestGenSchema(Schema):
|
|
source_code_being_tested: str
|
|
function_to_optimize: FunctionToOptimize
|
|
helper_function_names: list[str] | None = None # This is the only one we should use
|
|
dependent_function_names: list[str] | None = None # Only for backwards compatibility
|
|
module_path: str
|
|
test_module_path: str
|
|
test_framework: str # "pytest", "unittest", "jest", "mocha"
|
|
test_timeout: int
|
|
trace_id: str
|
|
python_version: str | None = None # Made optional for multi-language support
|
|
language: str = "python" # NEW: language identifier (python, javascript, typescript)
|
|
language_version: str | None = None # NEW: e.g., "ES2022", "Node 20", or Python version
|
|
codeflash_version: str | None = None
|
|
test_index: int | None = None
|
|
is_async: bool | None = False
|
|
call_sequence: int | None = None
|
|
|
|
@model_validator(mode="after")
|
|
def helper_function_names_validator(self) -> Self:
|
|
# To maintain backwards compatibility
|
|
if self.dependent_function_names is None and self.helper_function_names is None:
|
|
raise ValueError("either field 'helper_function_names' or 'dependent_function_names' is required")
|
|
if self.helper_function_names is not None:
|
|
return self
|
|
self.helper_function_names = self.dependent_function_names
|
|
self.dependent_function_names = None
|
|
return self
|
|
|
|
|
|
class TestGenResponseSchema(Schema):
|
|
generated_tests: str
|
|
instrumented_behavior_tests: str
|
|
instrumented_perf_tests: str
|
|
|
|
|
|
class TestGenErrorResponseSchema(Schema):
|
|
error: str
|
|
|
|
|
|
class TestGenerationFailedError(Exception):
|
|
pass
|