221 lines
8.6 KiB
Go
221 lines
8.6 KiB
Go
|
package workbench_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"slices"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesscontrol"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/workbench"
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
func accountScopeFixture(t *testing.T) (*gorm.DB, *gin.Engine) {
|
||
|
|
t.Helper()
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
db := operationTestDB(t)
|
||
|
|
if err := db.AutoMigrate(&types.User{}, &types.TenantMember{}, &accesscontrol.Assignment{}, &workbench.Balance{}, &workbench.Transaction{}, &workbench.UsageQuota{}, &workbench.Usage{}, &workbench.BalanceHold{}); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
deleted := gorm.DeletedAt{Time: time.Now().Add(-time.Hour), Valid: true}
|
||
|
|
for _, row := range []struct {
|
||
|
|
id string
|
||
|
|
tenant uint64
|
||
|
|
role types.TenantRole
|
||
|
|
status types.TenantMemberStatus
|
||
|
|
deletedUser bool
|
||
|
|
deletedMember bool
|
||
|
|
}{
|
||
|
|
{"admin", 1, types.TenantRoleAdmin, types.TenantMemberStatusActive, false, false},
|
||
|
|
{"ordinary", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, false, false},
|
||
|
|
{"inactive", 1, types.TenantRoleViewer, types.TenantMemberStatusActive, false, false},
|
||
|
|
{"suspended", 1, types.TenantRoleContributor, types.TenantMemberStatusSuspended, false, false},
|
||
|
|
{"acceptance-legitimate", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, false, false},
|
||
|
|
{"reenrolled", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, false, false},
|
||
|
|
{"deleted-user", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, true, false},
|
||
|
|
{"deleted-member", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, false, true},
|
||
|
|
{"deleted-both", 1, types.TenantRoleContributor, types.TenantMemberStatusActive, true, true},
|
||
|
|
{"other-tenant", 2, types.TenantRoleContributor, types.TenantMemberStatusActive, false, false},
|
||
|
|
} {
|
||
|
|
user := types.User{ID: row.id, Username: row.id, Email: row.id + "@example.test", TenantID: row.tenant, IsSystemAdmin: row.id == "admin"}
|
||
|
|
member := types.TenantMember{TenantID: row.tenant, UserID: row.id, Role: row.role, Status: row.status}
|
||
|
|
if row.deletedUser {
|
||
|
|
user.DeletedAt = deleted
|
||
|
|
}
|
||
|
|
if row.deletedMember {
|
||
|
|
member.DeletedAt = deleted
|
||
|
|
}
|
||
|
|
if err := db.Create(&user).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := db.Create(&member).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := db.Create(&workbench.UsageQuota{ID: row.id, TenantID: row.tenant, UserID: row.id, Period: "daily", OveragePolicy: "warn"}).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := db.Model(&types.User{}).Where("id=?", "inactive").Update("is_active", false).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := db.Create(&types.TenantMember{TenantID: 1, UserID: "reenrolled", Role: types.TenantRoleViewer, DeletedAt: deleted}).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := db.Create(&types.TenantMember{TenantID: 1, UserID: "missing-user", Role: types.TenantRoleContributor}).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
// A balance in another tenant must not be displayed for this tenant's account.
|
||
|
|
for _, balance := range []workbench.Balance{{TenantID: 1, UserID: "ordinary", Balance: 70, Reserved: 20}, {TenantID: 2, UserID: "ordinary", Balance: 900}} {
|
||
|
|
if err := db.Create(&balance).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
h := workbench.NewHandler(workbench.NewService(db))
|
||
|
|
router := gin.New()
|
||
|
|
router.GET("/accounts", h.Accounts)
|
||
|
|
router.POST("/accounts/grant", h.GrantBalance)
|
||
|
|
router.GET("/quotas", h.Quotas)
|
||
|
|
router.PUT("/quotas/:userID", h.SaveQuota)
|
||
|
|
router.GET("/users", accesscontrol.NewService(db).ListUsers)
|
||
|
|
return db, router
|
||
|
|
}
|
||
|
|
|
||
|
|
func accountScopeRequest(t *testing.T, router *gin.Engine, tenant uint64, role types.TenantRole, method, path string, body any) *httptest.ResponseRecorder {
|
||
|
|
t.Helper()
|
||
|
|
ctx := context.WithValue(operationTestContext(), types.TenantIDContextKey, tenant)
|
||
|
|
ctx = context.WithValue(ctx, types.TenantRoleContextKey, role)
|
||
|
|
encoded, err := json.Marshal(body)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
request := httptest.NewRequest(method, path, bytes.NewReader(encoded)).WithContext(ctx)
|
||
|
|
request.Header.Set("Content-Type", "application/json")
|
||
|
|
response := httptest.NewRecorder()
|
||
|
|
router.ServeHTTP(response, request)
|
||
|
|
return response
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBalanceAccountsMatchAccessControlLiveUsers(t *testing.T) {
|
||
|
|
_, router := accountScopeFixture(t)
|
||
|
|
expected := []string{"acceptance-legitimate", "admin", "inactive", "ordinary", "reenrolled", "suspended"}
|
||
|
|
for _, tenant := range []uint64{1, 2, 3} {
|
||
|
|
want := expected
|
||
|
|
if tenant == 2 {
|
||
|
|
want = []string{"other-tenant"}
|
||
|
|
} else if tenant == 3 {
|
||
|
|
want = []string{}
|
||
|
|
}
|
||
|
|
for _, path := range []string{"/users", "/accounts", "/quotas"} {
|
||
|
|
response := accountScopeRequest(t, router, tenant, types.TenantRoleAdmin, http.MethodGet, path, nil)
|
||
|
|
if response.Code != http.StatusOK {
|
||
|
|
t.Fatalf("%s tenant %d: %d %s", path, tenant, response.Code, response.Body.String())
|
||
|
|
}
|
||
|
|
var payload struct {
|
||
|
|
Data []struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
UserID string `json:"user_id"`
|
||
|
|
Balance int64 `json:"balance"`
|
||
|
|
Reserved int64 `json:"reserved"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(response.Body.Bytes(), &payload); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if payload.Data == nil {
|
||
|
|
t.Fatalf("%s must return an empty array instead of null", path)
|
||
|
|
}
|
||
|
|
ids := []string{}
|
||
|
|
for _, row := range payload.Data {
|
||
|
|
id := row.UserID
|
||
|
|
if path == "/users" {
|
||
|
|
id = row.ID
|
||
|
|
}
|
||
|
|
ids = append(ids, id)
|
||
|
|
if path == "/accounts" {
|
||
|
|
balance, reserved := int64(0), int64(0)
|
||
|
|
if tenant == 1 && id == "ordinary" {
|
||
|
|
balance, reserved = 70, 20
|
||
|
|
}
|
||
|
|
if row.Balance != balance || row.Reserved != reserved {
|
||
|
|
t.Fatalf("incorrect balance for %s in tenant %d: %+v", id, tenant, row)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if !slices.Equal(ids, want) {
|
||
|
|
t.Fatalf("%s tenant %d users = %v, want %v", path, tenant, ids, want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBalanceManagementRejectsDeletedAndOutOfScopeRecipients(t *testing.T) {
|
||
|
|
db, router := accountScopeFixture(t)
|
||
|
|
for _, user := range []string{"deleted-user", "deleted-member", "deleted-both", "missing-user", "other-tenant", "unknown"} {
|
||
|
|
grant := map[string]any{"user_id": user, "amount": 10, "reason": "scope regression", "idempotency_key": "scope-test-" + user}
|
||
|
|
quota := map[string]any{"period": "monthly", "enabled": true, "token_limit": 100, "overage_policy": "warn"}
|
||
|
|
for _, req := range []struct {
|
||
|
|
method, path string
|
||
|
|
body any
|
||
|
|
}{{http.MethodPost, "/accounts/grant", grant}, {http.MethodPut, "/quotas/" + user, quota}} {
|
||
|
|
response := accountScopeRequest(t, router, 1, types.TenantRoleAdmin, req.method, req.path, req.body)
|
||
|
|
if response.Code != http.StatusNotFound {
|
||
|
|
t.Fatalf("%s %s: got %d %s", req.method, req.path, response.Code, response.Body.String())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, check := range []struct {
|
||
|
|
model any
|
||
|
|
want int64
|
||
|
|
}{{&workbench.Transaction{}, 0}, {&workbench.Balance{}, 2}, {&workbench.UsageQuota{}, 10}} {
|
||
|
|
var count int64
|
||
|
|
if err := db.Model(check.model).Count(&count).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if count != check.want {
|
||
|
|
t.Fatalf("rejected request mutated %T: count=%d want=%d", check.model, count, check.want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Preserve ordinary grant behavior and idempotency for an existing member.
|
||
|
|
grant := map[string]any{"user_id": "ordinary", "amount": 10, "reason": "valid grant", "idempotency_key": "valid-grant-key"}
|
||
|
|
for range 2 {
|
||
|
|
response := accountScopeRequest(t, router, 1, types.TenantRoleAdmin, http.MethodPost, "/accounts/grant", grant)
|
||
|
|
if response.Code != http.StatusOK {
|
||
|
|
t.Fatalf("valid grant: %d %s", response.Code, response.Body.String())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
var balance workbench.Balance
|
||
|
|
if err := db.Where("tenant_id=1 AND user_id=?", "ordinary").Take(&balance).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if balance.Balance != 80 || balance.Reserved != 20 {
|
||
|
|
t.Fatalf("grant changed balance incorrectly: %+v", balance)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBalanceManagementRemainsAdminOnly(t *testing.T) {
|
||
|
|
_, router := accountScopeFixture(t)
|
||
|
|
for _, role := range []types.TenantRole{types.TenantRoleViewer, types.TenantRoleContributor} {
|
||
|
|
for _, req := range []struct {
|
||
|
|
method, path string
|
||
|
|
body any
|
||
|
|
}{
|
||
|
|
{http.MethodGet, "/accounts", nil},
|
||
|
|
{http.MethodGet, "/quotas", nil},
|
||
|
|
{http.MethodPost, "/accounts/grant", map[string]any{"user_id": "ordinary", "amount": 10, "idempotency_key": "denied-grant-key"}},
|
||
|
|
{http.MethodPut, "/quotas/ordinary", map[string]any{"period": "daily", "overage_policy": "warn"}},
|
||
|
|
} {
|
||
|
|
response := accountScopeRequest(t, router, 1, role, req.method, req.path, req.body)
|
||
|
|
if response.Code != http.StatusForbidden {
|
||
|
|
t.Fatalf("%s %s with role %s: got %d", req.method, req.path, role, response.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|