refactor: improve Vertex AI model conversion and add tests

- Rename compute_mode -> convert_model_for_vertex for clarity
- Move `import re` to module top-level (stdlib convention)
- Use greedy regex quantifier for more readable pattern matching
- Restore PEP 8 double blank line between top-level definitions
- Add test_client.py with 10 unit tests covering:
  - Vertex disabled (env unset, "0", empty)
  - Standard conversions (Opus, Sonnet, Haiku)
  - Edge cases (already-converted, non-Claude, no date suffix, empty)

Follow-up improvements from PR #129 review.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Auto
2026-01-30 10:20:03 +02:00
parent 813fcde18b
commit 79d02a1410
2 changed files with 110 additions and 4 deletions

View File

@@ -7,6 +7,7 @@ Functions for creating and configuring the Claude Agent SDK client.
import json import json
import os import os
import re
import shutil import shutil
import sys import sys
from pathlib import Path from pathlib import Path
@@ -68,7 +69,7 @@ EXTRA_READ_PATHS_BLOCKLIST = {
".netrc", ".netrc",
} }
def compute_mode(model: str) -> str: def convert_model_for_vertex(model: str) -> str:
""" """
Convert model name format for Vertex AI compatibility. Convert model name format for Vertex AI compatibility.
@@ -89,8 +90,7 @@ def compute_mode(model: str) -> str:
# Pattern: claude-{name}-{version}-{date} -> claude-{name}-{version}@{date} # Pattern: claude-{name}-{version}-{date} -> claude-{name}-{version}@{date}
# Example: claude-opus-4-5-20251101 -> claude-opus-4-5@20251101 # Example: claude-opus-4-5-20251101 -> claude-opus-4-5@20251101
# The date is always 8 digits at the end # The date is always 8 digits at the end
import re match = re.match(r'^(claude-.+)-(\d{8})$', model)
match = re.match(r'^(claude-[a-z0-9-]+?)-(\d{8})$', model)
if match: if match:
base_name, date = match.groups() base_name, date = match.groups()
return f"{base_name}@{date}" return f"{base_name}@{date}"
@@ -208,6 +208,7 @@ def get_extra_read_paths() -> list[Path]:
return validated_paths return validated_paths
# Feature MCP tools for feature/test management # Feature MCP tools for feature/test management
FEATURE_MCP_TOOLS = [ FEATURE_MCP_TOOLS = [
# Core feature operations # Core feature operations
@@ -438,7 +439,7 @@ def create_client(
is_vertex = sdk_env.get("CLAUDE_CODE_USE_VERTEX") == "1" is_vertex = sdk_env.get("CLAUDE_CODE_USE_VERTEX") == "1"
is_alternative_api = bool(base_url) or is_vertex is_alternative_api = bool(base_url) or is_vertex
is_ollama = "localhost:11434" in base_url or "127.0.0.1:11434" in base_url is_ollama = "localhost:11434" in base_url or "127.0.0.1:11434" in base_url
model = compute_mode(model) model = convert_model_for_vertex(model)
if sdk_env: if sdk_env:
print(f" - API overrides: {', '.join(sdk_env.keys())}") print(f" - API overrides: {', '.join(sdk_env.keys())}")
if is_vertex: if is_vertex:

105
test_client.py Normal file
View File

@@ -0,0 +1,105 @@
#!/usr/bin/env python3
"""
Client Utility Tests
====================
Tests for the client module utility functions.
Run with: python test_client.py
"""
import os
import unittest
from client import convert_model_for_vertex
class TestConvertModelForVertex(unittest.TestCase):
"""Tests for convert_model_for_vertex function."""
def setUp(self):
"""Save original env state."""
self._orig_vertex = os.environ.get("CLAUDE_CODE_USE_VERTEX")
def tearDown(self):
"""Restore original env state."""
if self._orig_vertex is None:
os.environ.pop("CLAUDE_CODE_USE_VERTEX", None)
else:
os.environ["CLAUDE_CODE_USE_VERTEX"] = self._orig_vertex
# --- Vertex AI disabled (default) ---
def test_returns_model_unchanged_when_vertex_disabled(self):
os.environ.pop("CLAUDE_CODE_USE_VERTEX", None)
self.assertEqual(
convert_model_for_vertex("claude-opus-4-5-20251101"),
"claude-opus-4-5-20251101",
)
def test_returns_model_unchanged_when_vertex_set_to_zero(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "0"
self.assertEqual(
convert_model_for_vertex("claude-opus-4-5-20251101"),
"claude-opus-4-5-20251101",
)
def test_returns_model_unchanged_when_vertex_set_to_empty(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = ""
self.assertEqual(
convert_model_for_vertex("claude-sonnet-4-5-20250929"),
"claude-sonnet-4-5-20250929",
)
# --- Vertex AI enabled: standard conversions ---
def test_converts_opus_model(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("claude-opus-4-5-20251101"),
"claude-opus-4-5@20251101",
)
def test_converts_sonnet_model(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("claude-sonnet-4-5-20250929"),
"claude-sonnet-4-5@20250929",
)
def test_converts_haiku_model(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("claude-3-5-haiku-20241022"),
"claude-3-5-haiku@20241022",
)
# --- Vertex AI enabled: already converted or non-matching ---
def test_already_vertex_format_unchanged(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("claude-opus-4-5@20251101"),
"claude-opus-4-5@20251101",
)
def test_non_claude_model_unchanged(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("gpt-4o"),
"gpt-4o",
)
def test_model_without_date_suffix_unchanged(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(
convert_model_for_vertex("claude-opus-4-5"),
"claude-opus-4-5",
)
def test_empty_string_unchanged(self):
os.environ["CLAUDE_CODE_USE_VERTEX"] = "1"
self.assertEqual(convert_model_for_vertex(""), "")
if __name__ == "__main__":
unittest.main()