httprunner/uixt/ai/ai_test.go

package ai

import (
	"context"
	"os"
	"testing"

	"github.com/httprunner/httprunner/v5/internal/builtin"
	"github.com/httprunner/httprunner/v5/uixt/option"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// hasRequiredEnvVars checks if the required environment variables are set for testing
func hasRequiredEnvVars() bool {
	// Check for OpenAI environment variables
	if os.Getenv("OPENAI_BASE_URL") != "" && os.Getenv("OPENAI_API_KEY") != "" {
		return true
	}
	// Check for GPT-4O specific environment variables
	if os.Getenv("OPENAI_GPT_4O_BASE_URL") != "" && os.Getenv("OPENAI_GPT_4O_API_KEY") != "" {
		return true
	}
	return false
}

func TestILLMServiceQuery(t *testing.T) {
	// Skip test if required environment variables are not set
	if !hasRequiredEnvVars() {
		t.Skip("Skipping test: required environment variables not set")
	}

	// Create LLM service
	service, err := NewLLMService(option.OPENAI_GPT_4O)
	require.NoError(t, err)
	require.NotNil(t, service)

	// Load test image
	screenshot, size, err := builtin.LoadImage("testdata/llk_1.png")
	require.NoError(t, err)

	// Test basic query functionality
	t.Run("BasicQuery", func(t *testing.T) {
		opts := &QueryOptions{
			Query:      "请描述这张图片中的内容",
			Screenshot: screenshot,
			Size:       size,
		}

		result, err := service.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("Query result: %s", result.Content)
	})

	// Test custom schema query
	t.Run("CustomSchemaQuery", func(t *testing.T) {
		type GameInfo struct {
			Content string   `json:"content"`
			Thought string   `json:"thought"`
			Rows    int      `json:"rows"`
			Cols    int      `json:"cols"`
			Icons   []string `json:"icons"`
		}

		opts := &QueryOptions{
			Query:        "请分析这个连连看游戏界面，告诉我有多少行多少列，有哪些不同类型的图案",
			Screenshot:   screenshot,
			Size:         size,
			OutputSchema: GameInfo{},
		}

		result, err := service.Query(context.Background(), opts)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		assert.NotEmpty(t, result.Content)
		assert.NotEmpty(t, result.Thought)
		assert.NotNil(t, result.Data)

		// Verify type conversion
		if gameInfo, ok := result.Data.(*GameInfo); ok {
			assert.NotEmpty(t, gameInfo.Content)
			assert.NotEmpty(t, gameInfo.Thought)
			assert.Greater(t, gameInfo.Rows, 0)
			assert.Greater(t, gameInfo.Cols, 0)
			assert.NotEmpty(t, gameInfo.Icons)
			t.Logf("Game info: %+v", gameInfo)
		} else {
			t.Errorf("Expected *GameInfo, got %T", result.Data)
		}
	})
}

func TestILLMServiceIntegration(t *testing.T) {
	// Skip test if required environment variables are not set
	if !hasRequiredEnvVars() {
		t.Skip("Skipping test: required environment variables not set")
	}

	// Create LLM service
	service, err := NewLLMService(option.OPENAI_GPT_4O)
	require.NoError(t, err)
	require.NotNil(t, service)

	// Load test image
	screenshot, size, err := builtin.LoadImage("testdata/llk_1.png")
	require.NoError(t, err)

	ctx := context.Background()

	// Test that all three methods work
	t.Run("AllMethods", func(t *testing.T) {
		// Test Query
		queryOpts := &QueryOptions{
			Query:      "请分析这张图片",
			Screenshot: screenshot,
			Size:       size,
		}
		queryResult, err := service.Query(ctx, queryOpts)
		assert.NoError(t, err)
		assert.NotNil(t, queryResult)
		t.Logf("Query result: %s", queryResult.Content)

		// Test Assert
		assertOpts := &AssertOptions{
			Assertion:  "这是一个连连看游戏界面",
			Screenshot: screenshot,
			Size:       size,
		}
		assertResult, err := service.Assert(ctx, assertOpts)
		assert.NoError(t, err)
		assert.NotNil(t, assertResult)
		t.Logf("Assert result: pass=%v, thought=%s", assertResult.Pass, assertResult.Thought)

		// Note: Planning test would require proper user instruction and message setup
		// which is more complex, so we skip it in this integration test
	})
}