Files
httprunner/uixt/ai/querier_test.go

417 lines
13 KiB
Go

package ai
import (
"context"
"encoding/json"
"fmt"
"testing"
"github.com/httprunner/httprunner/v5/internal/builtin"
"github.com/httprunner/httprunner/v5/uixt/option"
"github.com/httprunner/httprunner/v5/uixt/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Test data structures
// GameInfo represents basic game information for testing
type GameInfo struct {
Content string `json:"content"` // Description
Thought string `json:"thought"` // Reasoning
Rows int `json:"rows"` // Number of rows
Cols int `json:"cols"` // Number of columns
Icons []string `json:"icons"` // List of icon types
TotalIcons int `json:"totalNumber"` // Total number of icons
}
// GameAnalysisResult represents comprehensive game analysis for testing
type GameAnalysisResult struct {
Content string `json:"content"` // Human-readable description
Thought string `json:"thought"` // AI reasoning process
GameType string `json:"gameType"` // Type of game detected
Dimensions Dimensions `json:"dimensions"` // Grid dimensions
Elements []Element `json:"elements"` // Game elements detected
Statistics Statistics `json:"statistics"` // Game statistics
}
type Dimensions struct {
Rows int `json:"rows"` // Number of rows
Cols int `json:"cols"` // Number of columns
}
type Element struct {
Type string `json:"type"` // Element type/name
Position Position `json:"position"` // Position in grid
BoundBox BoundingBox `json:"boundBox"` // Pixel coordinates
}
type Position struct {
Row int `json:"row"` // Row index (0-based)
Col int `json:"col"` // Column index (0-based)
}
type BoundingBox struct {
X int `json:"x"` // Left coordinate
Y int `json:"y"` // Top coordinate
Width int `json:"width"` // Width in pixels
Height int `json:"height"` // Height in pixels
}
type Statistics struct {
TotalElements int `json:"totalElements"` // Total number of elements
UniqueTypes int `json:"uniqueTypes"` // Number of unique element types
TypeCounts []TypeCount `json:"typeCounts"` // Count of each type
}
type TypeCount struct {
Type string `json:"type"` // Element type
Count int `json:"count"` // Number of occurrences
}
// Test helper functions
func setupTestQuerier(t *testing.T) *Querier {
ctx := context.Background()
modelConfig, err := GetModelConfig(option.DOUBAO_SEED_1_6_250615)
require.NoError(t, err)
querier, err := NewQuerier(ctx, modelConfig)
require.NoError(t, err)
return querier
}
func loadTestImage(t *testing.T, path string) (string, types.Size) {
screenshot, size, err := builtin.LoadImage(path)
require.NoError(t, err)
return screenshot, size
}
// Test functions
func TestParseQueryResult(t *testing.T) {
tests := []struct {
name string
content string
expected *QueryResult
}{
{
name: "valid JSON response",
content: `{"content": "extracted information", "thought": "analysis complete"}`,
expected: &QueryResult{
Content: "extracted information",
Thought: "analysis complete",
},
},
{
name: "JSON in markdown",
content: "```json\n{\"content\": \"data from markdown\", \"thought\": \"parsed from code block\"}\n```",
expected: &QueryResult{
Content: "data from markdown",
Thought: "parsed from code block",
},
},
{
name: "plain text response",
content: "This is just plain text without JSON structure",
expected: &QueryResult{
Content: "This is just plain text without JSON structure",
Thought: "Failed to parse as JSON, returning raw content",
},
},
{
name: "invalid JSON",
content: `{"content": "incomplete json", "missing_closing_brace": true`,
expected: &QueryResult{
Content: "incomplete json",
Thought: "Partial extraction from malformed response",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := parseQueryResult(tt.content)
assert.NoError(t, err)
assert.Equal(t, tt.expected.Content, result.Content)
assert.Equal(t, tt.expected.Thought, result.Thought)
})
}
}
func TestModel(t *testing.T) {
// Test different models
models := []option.LLMServiceType{
option.DOUBAO_SEED_1_6_250615,
option.DOUBAO_1_5_THINKING_VISION_PRO_250428,
option.DOUBAO_1_5_UI_TARS_250328,
option.OPENAI_GPT_4O,
}
for _, path := range []string{"testdata/llk_1.png", "testdata/llk_4.png"} {
for _, modelType := range models {
t.Run(string(modelType), func(t *testing.T) {
modelConfig, err := GetModelConfig(modelType)
require.NoError(t, err)
querier, err := NewQuerier(context.Background(), modelConfig)
require.NoError(t, err)
// Load test image
screenshot, size := loadTestImage(t, path)
// Test query
opts := &QueryOptions{
Query: "请分析这个连连看游戏界面,告诉我有多少行多少列,有哪些不同类型的图案,图案总数是多少",
Screenshot: screenshot,
Size: size,
OutputSchema: GameInfo{},
}
result1, err := querier.Query(context.Background(), opts)
assert.NoError(t, err)
gameInfo, ok := result1.Data.(*GameInfo)
assert.True(t, ok)
jsonData1, _ := json.Marshal(gameInfo)
fmt.Printf("modelType: %v, gameInfo: %s\n", modelType, string(jsonData1))
opts2 := &QueryOptions{
Query: `Analyze this game interface and provide structured information about:
1. The type of game
2. Grid dimensions (rows and columns)
3. All game elements with their positions and types
4. Statistics about element distribution`,
Screenshot: screenshot,
Size: size,
OutputSchema: GameAnalysisResult{},
}
result2, err := querier.Query(context.Background(), opts2)
assert.NoError(t, err)
// Verify structured data
gameAnalysisResult, ok := result2.Data.(*GameAnalysisResult)
assert.True(t, ok)
jsonData2, _ := json.Marshal(gameAnalysisResult)
fmt.Printf("modelType: %v, gameAnalysisResult: %s\n", modelType, string(jsonData2))
opts3 := &QueryOptions{
Query: "给出第一个苹果的坐标",
Screenshot: screenshot,
Size: size,
OutputSchema: BoundingBox{},
}
result3, err := querier.Query(context.Background(), opts3)
assert.NoError(t, err)
boxInfo, ok := result3.Data.(*BoundingBox)
assert.True(t, ok)
jsonData3, _ := json.Marshal(boxInfo)
fmt.Printf("modelType: %v, thought: %v, boxInfo: %s\n",
modelType, result3.Thought, string(jsonData3))
})
}
}
}
// TestQueryFunctionality tests both basic and custom schema query functionality
func TestQueryFunctionality(t *testing.T) {
querier := setupTestQuerier(t)
screenshot, size := loadTestImage(t, "testdata/llk_1.png")
t.Run("BasicQuery", func(t *testing.T) {
opts := &QueryOptions{
Query: "这是一张连连看小游戏的界面,请分析游戏界面的基本信息",
Screenshot: screenshot,
Size: size,
}
result, err := querier.Query(context.Background(), opts)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Content)
assert.NotEmpty(t, result.Thought)
assert.Nil(t, result.Data) // Should be nil for standard query
t.Logf("Basic Query Result: %s", result.Content)
})
t.Run("CustomSchemaQuery", func(t *testing.T) {
opts := &QueryOptions{
Query: "请分析这个连连看游戏界面,告诉我有多少行多少列,有哪些不同类型的图案",
Screenshot: screenshot,
Size: size,
OutputSchema: GameInfo{},
}
result, err := querier.Query(context.Background(), opts)
assert.NoError(t, err)
// Verify structured data
gameInfo, ok := result.Data.(*GameInfo)
assert.True(t, ok)
assert.NotNil(t, gameInfo)
assert.NotEmpty(t, gameInfo.Content)
assert.NotEmpty(t, gameInfo.Thought)
assert.Equal(t, 4, gameInfo.Rows)
assert.Equal(t, 3, gameInfo.Cols)
assert.Equal(t, 5, gameInfo.TotalIcons)
t.Logf("Custom Schema Query Result: %+v", gameInfo)
})
t.Run("ComprehensiveAnalysis", func(t *testing.T) {
opts := &QueryOptions{
Query: `Analyze this game interface and provide structured information about:
1. The type of game
2. Grid dimensions (rows and columns)
3. All game elements with their positions and types
4. Statistics about element distribution`,
Screenshot: screenshot,
Size: size,
OutputSchema: GameAnalysisResult{},
}
result, err := querier.Query(context.Background(), opts)
assert.NoError(t, err)
// Verify structured data
gameAnalysisResult, ok := result.Data.(*GameAnalysisResult)
assert.True(t, ok)
assert.NotNil(t, gameAnalysisResult)
assert.NotEmpty(t, gameAnalysisResult.Content)
assert.NotEmpty(t, gameAnalysisResult.Thought)
assert.NotEmpty(t, gameAnalysisResult.GameType)
assert.Equal(t, 4, gameAnalysisResult.Dimensions.Rows)
assert.Equal(t, 3, gameAnalysisResult.Dimensions.Cols)
assert.Equal(t, 12, len(gameAnalysisResult.Elements))
t.Logf("Comprehensive Analysis Result: %+v", result.Data)
})
}
// TestQueryWithDifferentPrompts tests various types of queries on the same screenshot
func TestQueryWithDifferentPrompts(t *testing.T) {
querier := setupTestQuerier(t)
screenshot, size := loadTestImage(t, "testdata/llk_1.png")
queries := []string{
"请描述这张图片中的内容",
"这个游戏界面有多少行多少列?",
"请识别图片中所有不同类型的图案",
"请找出可以消除的图案对",
}
for i, query := range queries {
t.Run(fmt.Sprintf("Query_%d", i+1), func(t *testing.T) {
opts := &QueryOptions{
Query: query,
Screenshot: screenshot,
Size: size,
}
result, err := querier.Query(context.Background(), opts)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Content)
assert.NotEmpty(t, result.Thought)
t.Logf("Query %d: %s", i+1, query)
t.Logf("Answer: %s", result.Content)
})
}
}
// TestTypeConversionAndAssertion tests data type conversion and assertion functionality
func TestTypeConversionAndAssertion(t *testing.T) {
// Test data structure
type TestSchema struct {
Content string `json:"content"`
Thought string `json:"thought"`
Count int `json:"count"`
Items []string `json:"items"`
}
t.Run("ConvertQueryResultData", func(t *testing.T) {
// Create a QueryResult with structured data
testData := &TestSchema{
Content: "Test content",
Thought: "Test thought",
Count: 5,
Items: []string{"item1", "item2", "item3"},
}
result := &QueryResult{
Content: "Test content",
Thought: "Test thought",
Data: testData,
}
// Test type conversion
converted, err := ConvertQueryResultData[TestSchema](result)
assert.NoError(t, err)
assert.NotNil(t, converted)
assert.Equal(t, "Test content", converted.Content)
assert.Equal(t, "Test thought", converted.Thought)
assert.Equal(t, 5, converted.Count)
assert.Equal(t, []string{"item1", "item2", "item3"}, converted.Items)
})
t.Run("AutoTypeConversion", func(t *testing.T) {
// Simulate a JSON response from the model
jsonResponse := `{
"content": "Test content from model",
"thought": "Test reasoning process",
"count": 42,
"items": ["apple", "banana", "cherry"]
}`
// Test the parseCustomSchemaResult function directly
result, err := parseCustomSchemaResult(jsonResponse, TestSchema{})
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotNil(t, result.Data)
// Verify that Data is automatically converted to the correct type
typedData, ok := result.Data.(*TestSchema)
assert.True(t, ok, "Data should be automatically converted to *TestSchema")
assert.NotNil(t, typedData)
// Verify the content
assert.Equal(t, "Test content from model", typedData.Content)
assert.Equal(t, "Test reasoning process", typedData.Thought)
assert.Equal(t, 42, typedData.Count)
assert.Equal(t, []string{"apple", "banana", "cherry"}, typedData.Items)
// Verify that QueryResult fields are also populated
assert.Equal(t, "Test content from model", result.Content)
assert.Equal(t, "Test reasoning process", result.Thought)
})
t.Run("DirectTypeAssertion", func(t *testing.T) {
// Simulate a JSON response
jsonResponse := `{
"content": "Game analysis complete",
"thought": "Analyzed the game grid structure",
"count": 100,
"items": ["apple", "banana", "cherry", "grape"]
}`
// Test the parseCustomSchemaResult function
result, err := parseCustomSchemaResult(jsonResponse, TestSchema{})
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotNil(t, result.Data)
// Users can now directly use type assertion
if testData, ok := result.Data.(*TestSchema); ok {
assert.Equal(t, "Game analysis complete", testData.Content)
assert.Equal(t, "Analyzed the game grid structure", testData.Thought)
assert.Equal(t, 100, testData.Count)
assert.Equal(t, []string{"apple", "banana", "cherry", "grape"}, testData.Items)
} else {
t.Fatalf("Type assertion failed, Data type: %T", result.Data)
}
})
}