add some scriptsmain

1 files changed, 62 insertions, 0 deletions

diff --git a/internal/controlplane/limits_integration_test.go b/internal/controlplane/limits_integration_test.go
new file mode 100644
index 0000000..8087bdd
--- /dev/null
+++ b/internal/controlplane/limits_integration_test.go
@@ -0,0 +1,62 @@
+package controlplane
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+	"time"
+)
+
+func TestSetProjectLimitPostgres(t *testing.T) {
+	databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL")
+	if databaseURL == "" {
+		t.Skip("AIGW_TEST_DATABASE_URL is not set")
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	if err := MigrateDatabase(ctx, databaseURL); err != nil {
+		t.Fatal(err)
+	}
+	store, err := NewStore(ctx, Options{
+		DatabaseURL:   databaseURL,
+		CredentialKey: "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := store.Close(); err != nil {
+			t.Errorf("close control-plane store: %v", err)
+		}
+	})
+
+	suffix := time.Now().UnixNano()
+	tenant, _, err := store.CreateTenant(ctx, CreateTenantInput{Slug: fmt.Sprintf("limits-%d", suffix), Name: "Limits tenant"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if _, cleanupErr := store.db.Exec(context.Background(), `DELETE FROM tenants WHERE id=$1`, tenant.ID); cleanupErr != nil {
+			t.Errorf("cleanup limit integration data: %v", cleanupErr)
+		}
+	})
+	project, _, err := store.CreateProject(ctx, CreateProjectInput{TenantID: tenant.ID, Slug: "limits-project", Name: "Limits project"})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	got, _, err := store.SetProjectLimit(ctx, project.ID, SetProjectLimitInput{
+		RequestsPerMinute: 10, TokensPerMinute: 20, ConcurrentRequests: 3, MonthlySpendMicros: 40,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got.TenantID != tenant.ID || got.ProjectID != project.ID || got.TenantName != tenant.Name || got.ProjectName != project.Name {
+		t.Fatalf("unexpected limit identity: %+v", got)
+	}
+	if got.RequestsPerMinute != 10 || got.TokensPerMinute != 20 || got.Concurrent != 3 || got.MonthlySpendMicros != 40 {
+		t.Fatalf("unexpected limit values: %+v", got)
+	}
+}