|
|
|
@@ -0,0 +1,50 @@ |
|
|
|
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) |
|
|
|
} |
|
|
|
} |