- Add API server with WebSocket support and REST endpoints - Implement authentication system with API keys and permissions - Add task queue system with Redis backend and error handling - Include storage layer with database migrations and schemas - Add comprehensive logging, metrics, and telemetry - Implement security middleware and network utilities - Add experiment management and container orchestration - Include configuration management with smart defaults
117 lines
2.4 KiB
Go
117 lines
2.4 KiB
Go
package api
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/auth"
|
|
"github.com/jfraeys/fetch_ml/internal/queue"
|
|
)
|
|
|
|
func TestUserPermissions(t *testing.T) {
|
|
authConfig := &auth.AuthConfig{
|
|
Enabled: true,
|
|
APIKeys: map[auth.Username]auth.APIKeyEntry{
|
|
"admin": {
|
|
Hash: auth.APIKeyHash(auth.HashAPIKey("admin_key")),
|
|
Admin: true,
|
|
},
|
|
"scientist": {
|
|
Hash: auth.APIKeyHash(auth.HashAPIKey("ds_key")),
|
|
Admin: false,
|
|
Permissions: map[string]bool{
|
|
"jobs:create": true,
|
|
"jobs:read": true,
|
|
"jobs:update": true,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
apiKey string
|
|
permission string
|
|
want bool
|
|
}{
|
|
{"Admin can create", "admin_key", "jobs:create", true},
|
|
{"Scientist can create", "ds_key", "jobs:create", true},
|
|
{"Invalid key fails", "invalid_key", "jobs:create", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
user, err := authConfig.ValidateAPIKey(tt.apiKey)
|
|
|
|
if tt.apiKey == "invalid_key" {
|
|
if err == nil {
|
|
t.Error("Expected error for invalid API key")
|
|
}
|
|
return
|
|
}
|
|
|
|
if err != nil {
|
|
t.Errorf("Unexpected error: %v", err)
|
|
return
|
|
}
|
|
|
|
got := user.HasPermission(tt.permission)
|
|
if got != tt.want {
|
|
t.Errorf("HasPermission() = %v, want %v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestTaskOwnership(t *testing.T) {
|
|
tasks := []*queue.Task{
|
|
{
|
|
ID: "task1",
|
|
JobName: "user1_job",
|
|
UserID: "user1",
|
|
CreatedBy: "user1",
|
|
CreatedAt: time.Now(),
|
|
},
|
|
{
|
|
ID: "task2",
|
|
JobName: "user2_job",
|
|
UserID: "user2",
|
|
CreatedBy: "user2",
|
|
CreatedAt: time.Now(),
|
|
},
|
|
}
|
|
|
|
users := map[string]*auth.User{
|
|
"user1": {Name: "user1", Admin: false},
|
|
"user2": {Name: "user2", Admin: false},
|
|
"admin": {Name: "admin", Admin: true},
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
userName string
|
|
task *queue.Task
|
|
want bool
|
|
}{
|
|
{"User can view own task", "user1", tasks[0], true},
|
|
{"User cannot view other task", "user1", tasks[1], false},
|
|
{"Admin can view any task", "admin", tasks[1], true},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
user := users[tt.userName]
|
|
|
|
canAccess := false
|
|
if user.Admin {
|
|
canAccess = true
|
|
} else if tt.task.UserID == user.Name || tt.task.CreatedBy == user.Name {
|
|
canAccess = true
|
|
}
|
|
|
|
if canAccess != tt.want {
|
|
t.Errorf("Access = %v, want %v", canAccess, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|