112 lines
3.0 KiB
Python
112 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import patch
|
|
|
|
import yaml
|
|
|
|
from scripts.bezalel_gemma4_vps import (
|
|
build_deploy_mutation,
|
|
build_runpod_endpoint,
|
|
parse_deploy_response,
|
|
update_config_text,
|
|
verify_openai_chat,
|
|
)
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload: dict):
|
|
self._payload = json.dumps(payload).encode()
|
|
|
|
def read(self) -> bytes:
|
|
return self._payload
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
|
|
def test_build_deploy_mutation_uses_ollama_image_and_openai_port() -> None:
|
|
query = build_deploy_mutation(name="bezalel-gemma4", gpu_type="NVIDIA L40S", model_tag="gemma4:latest")
|
|
|
|
assert 'gpuTypeId: "NVIDIA L40S"' in query
|
|
assert 'imageName: "ollama/ollama:latest"' in query
|
|
assert 'ports: "11434/http"' in query
|
|
assert 'volumeMountPath: "/root/.ollama"' in query
|
|
|
|
|
|
def test_build_runpod_endpoint_appends_v1_suffix() -> None:
|
|
assert build_runpod_endpoint("abc123") == "https://abc123-11434.proxy.runpod.net/v1"
|
|
|
|
|
|
def test_parse_deploy_response_extracts_pod_id_and_endpoint() -> None:
|
|
payload = {
|
|
"data": {
|
|
"podFindAndDeployOnDemand": {
|
|
"id": "podxyz",
|
|
"desiredStatus": "RUNNING",
|
|
}
|
|
}
|
|
}
|
|
|
|
result = parse_deploy_response(payload)
|
|
|
|
assert result == {
|
|
"pod_id": "podxyz",
|
|
"desired_status": "RUNNING",
|
|
"base_url": "https://podxyz-11434.proxy.runpod.net/v1",
|
|
}
|
|
|
|
|
|
def test_update_config_text_upserts_big_brain_provider() -> None:
|
|
original = """
|
|
model:
|
|
default: kimi-k2.5
|
|
provider: kimi-coding
|
|
custom_providers:
|
|
- name: Big Brain
|
|
base_url: https://old-endpoint/v1
|
|
api_key: ''
|
|
model: gemma3:27b
|
|
"""
|
|
|
|
updated = update_config_text(original, base_url="https://new-pod-11434.proxy.runpod.net/v1", model="gemma4:latest")
|
|
parsed = yaml.safe_load(updated)
|
|
|
|
assert parsed["model"] == {"default": "kimi-k2.5", "provider": "kimi-coding"}
|
|
assert parsed["custom_providers"] == [
|
|
{
|
|
"name": "Big Brain",
|
|
"base_url": "https://new-pod-11434.proxy.runpod.net/v1",
|
|
"api_key": "",
|
|
"model": "gemma4:latest",
|
|
}
|
|
]
|
|
|
|
|
|
def test_verify_openai_chat_calls_chat_completions() -> None:
|
|
response_payload = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": "READY"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
with patch(
|
|
"scripts.bezalel_gemma4_vps.request.urlopen",
|
|
return_value=_FakeResponse(response_payload),
|
|
) as mocked:
|
|
result = verify_openai_chat("https://pod-11434.proxy.runpod.net/v1", model="gemma4:latest", prompt="say READY")
|
|
|
|
assert result == "READY"
|
|
req = mocked.call_args.args[0]
|
|
assert req.full_url == "https://pod-11434.proxy.runpod.net/v1/chat/completions"
|
|
payload = json.loads(req.data.decode())
|
|
assert payload["model"] == "gemma4:latest"
|
|
assert payload["messages"][0]["content"] == "say READY"
|