79 lines
2.6 KiB
Go
79 lines
2.6 KiB
Go
package runner
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
"github.com/riverqueue/river/rivertype"
|
|
)
|
|
|
|
func TestRiverJobStateCancellableOnlyAllowsLocalQueueStates(t *testing.T) {
|
|
cancellableStates := []rivertype.JobState{
|
|
rivertype.JobStateAvailable,
|
|
rivertype.JobStateScheduled,
|
|
rivertype.JobStateRetryable,
|
|
rivertype.JobStatePending,
|
|
}
|
|
for _, state := range cancellableStates {
|
|
if !riverJobStateCancellable(state) {
|
|
t.Fatalf("expected %s to be cancellable", state)
|
|
}
|
|
}
|
|
|
|
nonCancellableStates := []rivertype.JobState{
|
|
rivertype.JobStateRunning,
|
|
rivertype.JobStateCancelled,
|
|
rivertype.JobStateCompleted,
|
|
rivertype.JobStateDiscarded,
|
|
}
|
|
for _, state := range nonCancellableStates {
|
|
if riverJobStateCancellable(state) {
|
|
t.Fatalf("expected %s to be non-cancellable", state)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTaskAccessibleToUserHonorsAPIKeyAndGatewayUser(t *testing.T) {
|
|
task := store.GatewayTask{
|
|
UserID: "server-user",
|
|
GatewayUserID: "gateway-user",
|
|
APIKeyID: "api-key",
|
|
}
|
|
|
|
if !taskAccessibleToUser(task, &auth.User{GatewayUserID: "gateway-user"}) {
|
|
t.Fatal("gateway user should access own task")
|
|
}
|
|
if !taskAccessibleToUser(task, &auth.User{APIKeyID: "api-key", GatewayUserID: "gateway-user"}) {
|
|
t.Fatal("api key owner should access own task")
|
|
}
|
|
if taskAccessibleToUser(task, &auth.User{APIKeyID: "api-key", GatewayUserID: "other-user"}) {
|
|
t.Fatal("api key from another gateway user should not access task")
|
|
}
|
|
if taskAccessibleToUser(task, &auth.User{GatewayUserID: "other-user"}) {
|
|
t.Fatal("another gateway user should not access task")
|
|
}
|
|
|
|
serverMainTask := store.GatewayTask{
|
|
UserID: "server-user",
|
|
APIKeyID: "server-api-key",
|
|
}
|
|
if !taskAccessibleToUser(serverMainTask, &auth.User{ID: "server-user", Source: "server-main", APIKeyID: "server-api-key"}) {
|
|
t.Fatal("server-main api key should fall back to user id when gateway user id is absent")
|
|
}
|
|
if taskAccessibleToUser(serverMainTask, &auth.User{ID: "other-user", Source: "server-main", APIKeyID: "server-api-key"}) {
|
|
t.Fatal("server-main api key from another user should not access task")
|
|
}
|
|
}
|
|
|
|
func TestTaskCancelUnavailableReportsSubmittedFromRemoteTaskID(t *testing.T) {
|
|
result := taskCancelUnavailable(store.GatewayTask{ID: "task-1", RemoteTaskID: "remote-1"}, "不可取消")
|
|
|
|
if result.Cancelled || result.Cancellable {
|
|
t.Fatalf("unavailable result should not mark cancellation as done: %+v", result)
|
|
}
|
|
if !result.Submitted {
|
|
t.Fatalf("remote task id should report submitted: %+v", result)
|
|
}
|
|
}
|