feat(api): 支持取消本地排队任务

2026-06-09 19:09:20 +08:00
parent a8fa8dd212
commit e8df26da9b
6 changed files with 319 additions and 0 deletions
@@ -0,0 +1,78 @@
+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)
+	}
+}