yao/agent/api_test.go

package agent

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	httpTest "github.com/yaoapp/gou/http"
	"github.com/yaoapp/yao/config"
	"github.com/yaoapp/yao/helper"
	"github.com/yaoapp/yao/test"
)

func init() {
	// Set gin to release mode to reduce log output
	gin.SetMode(gin.ReleaseMode)
}

func TestAPI(t *testing.T) {
	// Disable test logging
	test.Prepare(t, config.Conf)
	defer test.Clean()

	// Redirect stdout to /dev/null
	oldStdout := os.Stdout
	null, _ := os.Open(os.DevNull)
	os.Stdout = null
	defer func() {
		os.Stdout = oldStdout
		null.Close()
	}()

	// test router
	router := testRouter(t)
	err := Agent.API(router, "/agent/chat")
	if err != nil {
		t.Fatal(err)
	}

	// test server
	host, shutdown := testServer(t, router)
	defer shutdown()

	tests := []struct {
		name       string
		url        string
		method     string
		headers    http.Header
		expectCode int
		expectBody string
	}{
		{
			name:       "Basic Chat Request",
			url:        fmt.Sprintf("/agent/chat?content=hello&token=%s", testToken()),
			method:     "GET",
			headers:    http.Header{"Content-Type": []string{"application/json"}},
			expectBody: `{`,
		},
		{
			name:       "Chat with System Message",
			url:        fmt.Sprintf("/agent/chat?content=hello&system=You are a helpful assistant&token=%s", testToken()),
			method:     "GET",
			headers:    http.Header{"Content-Type": []string{"application/json"}},
			expectBody: `{`,
		},
		{
			name:       "Chat with Model Parameter",
			url:        fmt.Sprintf("/agent/chat?content=hello&model=gpt-3.5-turbo&token=%s", testToken()),
			method:     "GET",
			headers:    http.Header{"Content-Type": []string{"application/json"}},
			expectBody: `{`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			url := fmt.Sprintf("%s%s", host, tt.url)
			res := []byte{}
			req := httpTest.New(url).WithHeader(tt.headers)

			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			req.Stream(ctx, tt.method, nil, func(data []byte) int {
				res = append(res, data...)
				return 1
			})

			assert.Contains(t, string(res), tt.expectBody)
		})
	}
}

func TestAPIAuth(t *testing.T) {
	test.Prepare(t, config.Conf)
	defer test.Clean()

	// Redirect stdout and stderr to /dev/null
	oldStdout := os.Stdout
	oldStderr := os.Stderr
	null, _ := os.Open(os.DevNull)
	os.Stdout = null
	os.Stderr = null
	defer func() {
		os.Stdout = oldStdout
		os.Stderr = oldStderr
		null.Close()
	}()

	router := testRouter(t)
	err := Agent.API(router, "/agent/chat")
	if err != nil {
		t.Fatal(err)
	}

	// Separate tests for authentication errors and parameter validation errors
	authTests := []struct {
		name       string
		url        string
		method     string
		expectCode int
	}{
		{
			name:       "Missing Token",
			url:        "/agent/chat?content=hello",
			method:     "GET",
			expectCode: http.StatusUnauthorized,
		},
		{
			name:       "Invalid Token",
			url:        "/agent/chat?content=hello&token=invalid",
			method:     "GET",
			expectCode: http.StatusUnauthorized,
		},
	}

	// Test authentication errors (will panic)
	for _, tt := range authTests {
		t.Run(tt.name, func(t *testing.T) {
			response := httptest.NewRecorder()
			req, _ := http.NewRequest(tt.method, tt.url, nil)
			assert.Panics(t, func() {
				router.ServeHTTP(response, req)
			})
		})
	}

	// Test parameter validation errors (will return status code)
	validationTests := []struct {
		name       string
		url        string
		method     string
		expectCode int
	}{
		{
			name:       "Missing Content",
			url:        fmt.Sprintf("/agent/chat?token=%s", testToken()),
			method:     "GET",
			expectCode: http.StatusBadRequest,
		},
	}

	// Test parameter validation errors (return status code)
	for _, tt := range validationTests {
		t.Run(tt.name, func(t *testing.T) {
			response := httptest.NewRecorder()
			req, _ := http.NewRequest(tt.method, tt.url, nil)
			router.ServeHTTP(response, req)
			assert.Equal(t, tt.expectCode, response.Code)
		})
	}
}

// Helper functions
func testServer(t *testing.T, router *gin.Engine) (string, func()) {
	l, err := net.Listen("tcp4", ":0")
	if err != nil {
		t.Fatal(err)
	}

	srv := &http.Server{Addr: ":0", Handler: router}

	go func() {
		if err := srv.Serve(l); err != nil && err != http.ErrServerClosed {
			return
		}
	}()

	addr := strings.Split(l.Addr().String(), ":")
	if len(addr) != 2 {
		t.Fatal("invalid address")
	}

	host := fmt.Sprintf("http://127.0.0.1:%s", addr[1])
	time.Sleep(50 * time.Millisecond)

	shutdown := func() {
		srv.Close()
		l.Close()
	}
	return host, shutdown
}

func testRouter(t *testing.T) *gin.Engine {
	err := Load(config.Conf)
	if err != nil {
		t.Fatal(err)
	}

	router := gin.New() // Use gin.New() instead of gin.Default() to avoid default logging middleware
	return router
}

func testToken() string {
	token := helper.JwtMake(1,
		map[string]interface{}{
			"id":   1,
			"name": "Test",
		},
		map[string]interface{}{
			"exp": 3600,
			"sid": "123456",
		})
	return token.Token
}