new-api/relay/channel/openai/upstream_body_test.go

package openai

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/dto"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestResponsesStreamEventErrorParsing(t *testing.T) {
	t.Run("standalone error event", func(t *testing.T) {
		data := `{"type":"error","error":{"type":"too_many_requests","code":"too_many_requests","message":"Too Many Requests","param":null}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "error", streamResp.Type)
		assert.NotNil(t, streamResp.Error)

		oaiErr := dto.GetOpenAIError(streamResp.Error)
		require.NotNil(t, oaiErr)
		assert.Equal(t, "too_many_requests", oaiErr.Type)
		assert.Equal(t, "Too Many Requests", oaiErr.Message)
	})

	t.Run("response.failed event", func(t *testing.T) {
		data := `{"type":"response.failed","response":{"id":"test-id","status":"failed","error":{"type":"server_error","message":"Internal server error"}}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.failed", streamResp.Type)
		require.NotNil(t, streamResp.Response)

		oaiErr := streamResp.Response.GetOpenAIError()
		require.NotNil(t, oaiErr)
		assert.Equal(t, "server_error", oaiErr.Type)
	})

	t.Run("response.error event with error in response object", func(t *testing.T) {
		data := `{"type":"response.error","response":{"id":"test-id","status":"failed","error":{"type":"invalid_request_error","message":"Invalid model"}}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.error", streamResp.Type)

		oaiErr := streamResp.Response.GetOpenAIError()
		require.NotNil(t, oaiErr)
		assert.Equal(t, "invalid_request_error", oaiErr.Type)
	})

	t.Run("normal event has no error", func(t *testing.T) {
		data := `{"type":"response.output_text.delta","delta":"hello"}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.output_text.delta", streamResp.Type)
		assert.Nil(t, streamResp.Error)
	})
}

func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(`{"incomplete":`)),
	}

	usage, err := OaiResponsesHandler(c, &relaycommon.RelayInfo{}, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.NotEmpty(t, err.UpstreamBody)
	assert.Contains(t, err.UpstreamBody, `{"incomplete":`)
}

func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(`{"incomplete":`)),
	}

	usage, err := OaiResponsesToChatHandler(c, &relaycommon.RelayInfo{}, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.NotEmpty(t, err.UpstreamBody)
	assert.Contains(t, err.UpstreamBody, `{"incomplete":`)
}