new-api/relay/channel/task/doubao/adaptor_test.go

package doubao

import (
	"testing"

	"github.com/QuantumNous/new-api/model"
	"github.com/stretchr/testify/require"
)

func TestParseTaskResult_SuccessWithUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	taskInfo, err := adaptor.ParseTaskResult([]byte(`{
		"status":"succeeded",
		"content":{"video_url":"https://example.test/video.mp4"},
		"usage":{"completion_tokens":12,"total_tokens":3456}
	}`))

	require.NoError(t, err)
	require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status)
	require.Equal(t, "https://example.test/video.mp4", taskInfo.Url)
	require.Equal(t, 12, taskInfo.CompletionTokens)
	require.Equal(t, 3456, taskInfo.TotalTokens)
}

func TestParseTaskResult_SuccessWithoutUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	taskInfo, err := adaptor.ParseTaskResult([]byte(`{
		"status":"succeeded",
		"content":{"video_url":"https://example.test/video.mp4"}
	}`))

	require.NoError(t, err)
	require.Equal(t, string(model.TaskStatusSuccess), taskInfo.Status)
	require.Zero(t, taskInfo.TotalTokens)
}

func TestParseTaskResult_NonSuccessStatusesDoNotRequireUsage(t *testing.T) {
	adaptor := &TaskAdaptor{}
	cases := map[string]model.TaskStatus{
		"queued":     model.TaskStatusQueued,
		"processing": model.TaskStatusInProgress,
		"failed":     model.TaskStatusFailure,
	}
	for status, want := range cases {
		taskInfo, err := adaptor.ParseTaskResult([]byte(`{"status":"` + status + `"}`))
		require.NoError(t, err)
		require.Equal(t, string(want), taskInfo.Status)
		require.Zero(t, taskInfo.TotalTokens)
	}
}