ai-gateway-go/internal/portal/portal_integration_test.go

package portal

import (
	"bytes"
	"context"
	"encoding/base64"
	"net/http"
	"os"
	"testing"

	"aigateway.local/core/internal/apikey"
	"aigateway.local/core/internal/identity"
	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/cryptox"
	"aigateway.local/core/internal/platform/database"
	"aigateway.local/core/internal/workbench"
)

func TestPortalWorkspacePostgreSQLLifecycle(t *testing.T) {
	databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
	if databaseURL == "" {
		t.Skip("WORKBENCH_TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8})
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	const actorID = "11111111-1111-4111-8111-111111111111"
	const userID = "22222222-2222-4222-8222-222222222222"
	_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'portal-test-admin','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
	if err != nil {
		t.Fatal(err)
	}
	_, err = pool.Exec(ctx, `INSERT INTO gateway.portal_users(id,account,name,role,permissions,password_hash,auth_source,active) VALUES($1,'portal-test-user','Portal Test','member','{}','test','local',true) ON CONFLICT(id) DO NOTHING`, userID)
	if err != nil {
		t.Fatal(err)
	}
	cleanup := func() {
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.api_keys WHERE id IN (SELECT api_key_id FROM gateway.application_runtime_credentials c JOIN gateway.applications a ON a.id=c.application_id WHERE a.code='portal_test_app')`)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.model_access_requests WHERE portal_user_id=$1`, userID)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.prompt_favorites WHERE portal_user_id=$1`, userID)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.applications WHERE code='portal_test_app'`)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.prompt_templates WHERE name='portal-test-prompt'`)
		_, _ = pool.Exec(ctx, `DELETE FROM gateway.portal_users WHERE id=$1`, userID)
	}
	cleanup()
	_, err = pool.Exec(ctx, `INSERT INTO gateway.portal_users(id,account,name,role,permissions,password_hash,auth_source,active) VALUES($1,'portal-test-user','Portal Test','member','{}','test','local',true)`, userID)
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()

	assets := workbench.NewService(pool)
	prompt, err := assets.CreatePrompt(ctx, workbench.PromptInput{Name: "portal-test-prompt", Enabled: true, Content: "hello"}, actorID)
	if err != nil {
		t.Fatal(err)
	}
	app, err := assets.SaveApplication(ctx, workbench.Application{Code: "portal_test_app", Name: "Portal Test", Status: "draft", DraftConfig: workbench.ApplicationConfig{Model: "test-model", PromptTemplateID: prompt.ID}}, actorID, true)
	if err != nil {
		t.Fatal(err)
	}
	if _, err = assets.PublishApplication(ctx, app.ID, "portal test", actorID); err != nil {
		t.Fatal(err)
	}
	key := base64.StdEncoding.EncodeToString(make([]byte, 32))
	cipher, err := cryptox.NewKeyring(key, 1, "", "portal-test-tool")
	if err != nil {
		t.Fatal(err)
	}
	service := NewService(pool, assets, workbench.NewToolService(assets, cipher, true), nil)
	runtime := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !bytes.HasPrefix([]byte(r.Header.Get("Authorization")), []byte("Bearer gw_")) {
			t.Errorf("runtime credential was not supplied")
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"portal answer"}}]}`))
	})
	service.SetApplicationRuntime(NewRuntimeCredentials(pool, apikey.NewRepository(pool), cipher), runtime)
	account := identity.Account{ID: userID, Kind: identity.KindPortal}
	apps, err := service.Applications(ctx, account)
	if err != nil || len(apps) != 1 {
		t.Fatalf("apps=%#v err=%v", apps, err)
	}
	prompts, err := service.Prompts(ctx, account)
	if err != nil || len(prompts) != 1 || prompts[0].Favorite {
		t.Fatalf("prompts=%#v err=%v", prompts, err)
	}
	if err = service.SetFavorite(ctx, account, prompt.ID, true); err != nil {
		t.Fatal(err)
	}
	promptView, err := service.Prompt(ctx, account, prompt.ID)
	if err != nil || !promptView.Favorite {
		t.Fatalf("prompt=%#v err=%v", promptView, err)
	}
	request, err := service.CreateModelRequest(ctx, account, ModelRequest{Model: "test-model", Reason: "project", RequestedRPM: 10, RequestedMonthlyTokens: 1000})
	if err != nil || request.Status != "pending" {
		t.Fatalf("request=%#v err=%v", request, err)
	}
	request, err = service.DecideModelRequest(ctx, request.ID, "approved", "approved", actorID)
	if err != nil || request.Status != "approved" || request.DecidedAt == nil {
		t.Fatalf("decision=%#v err=%v", request, err)
	}
	conversation, err := service.CreateConversation(ctx, account, "portal_test_app")
	if err != nil {
		t.Fatal(err)
	}
	response, err := service.AppendConversationMessage(ctx, account, "portal_test_app", conversation.ID, "hello", nil)
	if err != nil || response["history_integrity"] != "verified" {
		t.Fatalf("response=%#v err=%v", response, err)
	}
	conversation, err = service.Conversation(ctx, account, "portal_test_app", conversation.ID)
	if err != nil || len(conversation.Messages) != 2 || conversation.Messages[1].Content != "portal answer" {
		t.Fatalf("conversation=%#v err=%v", conversation, err)
	}
}