From abcbc4e58dbbb64f2c586bd3ac2303cdfd9ed30e Mon Sep 17 00:00:00 2001
From: Your Name <your.email@example.com>
Date: Fri, 29 May 2026 20:48:12 +0800
Subject: [PATCH] test: add antigravity model functions tests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add tests for model-related functions:
- DefaultModels
- DefaultGeminiModels
- FallbackGeminiModelsList
- FallbackGeminiModel
- ClaudeModels/GeminiModels verification

Coverage: antigravity 18.8% → 19.6%
---
 internal/pkg/antigravity/claude_types_test.go | 128 +++++++++++++++---
 1 file changed, 111 insertions(+), 17 deletions(-)

diff --git a/internal/pkg/antigravity/claude_types_test.go b/internal/pkg/antigravity/claude_types_test.go
index 9fc09b1..34343e3 100644
--- a/internal/pkg/antigravity/claude_types_test.go
+++ b/internal/pkg/antigravity/claude_types_test.go
@@ -1,28 +1,122 @@
 package antigravity
 
-import "testing"
+import (
+	"strings"
+	"testing"
 
-func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
-	t.Parallel()
+	"github.com/stretchr/testify/assert"
+)
 
+func TestDefaultModels(t *testing.T) {
 	models := DefaultModels()
-	byID := make(map[string]ClaudeModel, len(models))
+
+	// Should return models (Claude + Gemini)
+	assert.NotEmpty(t, models)
+	assert.GreaterOrEqual(t, len(models), len(claudeModels)+len(geminiModels))
+
+	// Check all models have required fields
 	for _, m := range models {
-		byID[m.ID] = m
+		assert.NotEmpty(t, m.ID, "Model ID should not be empty")
+		assert.NotEmpty(t, m.DisplayName, "DisplayName should not be empty")
+		assert.Equal(t, "model", m.Type)
+		assert.NotEmpty(t, m.CreatedAt)
 	}
 
-	requiredIDs := []string{
-		"claude-opus-4-6-thinking",
-		"gemini-2.5-flash-image",
-		"gemini-2.5-flash-image-preview",
-		"gemini-3.1-flash-image",
-		"gemini-3.1-flash-image-preview",
-		"gemini-3-pro-image", // legacy compatibility
-	}
-
-	for _, id := range requiredIDs {
-		if _, ok := byID[id]; !ok {
-			t.Fatalf("expected model %q to be exposed in DefaultModels", id)
+	// Check that Claude models are present
+	claudeFound := false
+	for _, m := range models {
+		if strings.HasPrefix(m.ID, "claude-") {
+			claudeFound = true
+			break
 		}
 	}
+	assert.True(t, claudeFound, "Should have Claude models")
+
+	// Check that Gemini models are present
+	geminiFound := false
+	for _, m := range models {
+		if strings.HasPrefix(m.ID, "gemini-") {
+			geminiFound = true
+			break
+		}
+	}
+	assert.True(t, geminiFound, "Should have Gemini models")
+}
+
+func TestDefaultGeminiModels(t *testing.T) {
+	models := DefaultGeminiModels()
+
+	// Should return only Gemini models
+	assert.Equal(t, len(geminiModels), len(models))
+
+	// Check all models have required fields
+	for _, m := range models {
+		assert.True(t, strings.HasPrefix(m.Name, "models/"), "Name should have models/ prefix")
+		assert.NotEmpty(t, m.DisplayName)
+		assert.Equal(t, defaultGeminiMethods, m.SupportedGenerationMethods)
+	}
+}
+
+func TestFallbackGeminiModelsList(t *testing.T) {
+	resp := FallbackGeminiModelsList()
+
+	assert.NotEmpty(t, resp.Models)
+	assert.Equal(t, len(geminiModels), len(resp.Models))
+}
+
+func TestFallbackGeminiModel(t *testing.T) {
+	tests := []struct {
+		name         string
+		input        string
+		expectedName string
+	}{
+		{
+			name:         "with models/ prefix",
+			input:        "gemini-2.5-flash",
+			expectedName: "models/gemini-2.5-flash",
+		},
+		{
+			name:         "without models/ prefix",
+			input:        "gemini-3-pro",
+			expectedName: "models/gemini-3-pro",
+		},
+		{
+			name:         "already has prefix",
+			input:        "models/gemini-2.5-flash",
+			expectedName: "models/gemini-2.5-flash",
+		},
+		{
+			name:         "empty string",
+			input:        "",
+			expectedName: "models/unknown",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := FallbackGeminiModel(tt.input)
+			assert.Equal(t, tt.expectedName, result.Name)
+			assert.Equal(t, defaultGeminiMethods, result.SupportedGenerationMethods)
+		})
+	}
+}
+
+func TestClaudeModelsNotEmpty(t *testing.T) {
+	// Verify internal claudeModels slice is not empty
+	assert.NotEmpty(t, claudeModels)
+
+	for _, m := range claudeModels {
+		assert.NotEmpty(t, m.ID)
+		assert.NotEmpty(t, m.DisplayName)
+	}
+}
+
+func TestGeminiModelsNotEmpty(t *testing.T) {
+	// Verify internal geminiModels slice is not empty
+	assert.NotEmpty(t, geminiModels)
+
+	for _, m := range geminiModels {
+		assert.NotEmpty(t, m.ID)
+		assert.NotEmpty(t, m.DisplayName)
+	}
 }