Refactor large Rust integration tests

2026-06-21 07:12:33 +00:00
parent d01c8e1f1a
commit a31862aeeb
23 changed files with 6119 additions and 5284 deletions
@@ -0,0 +1,276 @@
+#![allow(dead_code, unused_imports)]
+
+use std::collections::BTreeMap;
+
+use crank_core::{
+    AgentId, AgentOperationBinding, AgentStatus, AgentVersion, ApiKeyHeaderAuthConfig, AuthConfig,
+    AuthKind, AuthProfile, ConfigExport, ExecutionConfig, ExportMode, GeneratedDraft,
+    GeneratedDraftStatus, HttpMethod, InvocationLog, MembershipRole, OperationId,
+    OperationSecurityLevel, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyScope,
+    PlatformApiKeyStatus, Protocol, RestTarget, RetryPolicy, Samples, SecretId, Target,
+    ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState, Workspace, WorkspaceId,
+};
+use crank_mapping::{MappingRule, MappingSet};
+use crank_schema::{Schema, SchemaKind};
+use serde_json::json;
+use sqlx::{Executor, PgPool, postgres::PgPoolOptions};
+use time::{OffsetDateTime, format_description::well_known::Rfc3339};
+
+use crank_registry::{
+    CreateAgentRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
+    CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
+    DescriptorMetadata, OperationSampleMetadata, PlatformApiKeyRecord, PostgresRegistry,
+    PublishAgentRequest, PublishRequest, RegistryError, RegistryOperation, SampleKind,
+    SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
+    WorkspaceRecord, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
+};
+
+fn test_workspace_id() -> WorkspaceId {
+    WorkspaceId::new("ws_default")
+}
+
+fn timestamp(value: &str) -> OffsetDateTime {
+    OffsetDateTime::parse(value, &Rfc3339).unwrap()
+}
+
+pub(super) fn test_operation(id: &str, version: u32, status: OperationStatus) -> RegistryOperation {
+    RegistryOperation {
+        id: OperationId::new(id),
+        name: format!("{id}_tool"),
+        display_name: format!("Display {id}"),
+        category: "general".to_owned(),
+        protocol: Protocol::Rest,
+        security_level: OperationSecurityLevel::Standard,
+        status,
+        version,
+        target: Target::Rest(RestTarget {
+            base_url: "https://api.example.com".to_owned(),
+            method: HttpMethod::Post,
+            path_template: "/v1/leads".to_owned(),
+            static_headers: BTreeMap::from([("X-Static".to_owned(), "true".to_owned())]),
+        }),
+        input_schema: Schema {
+            kind: SchemaKind::Object,
+            description: Some("input".to_owned()),
+            required: true,
+            nullable: false,
+            default_value: None,
+            fields: BTreeMap::from([(
+                "email".to_owned(),
+                Schema {
+                    kind: SchemaKind::String,
+                    description: None,
+                    required: true,
+                    nullable: false,
+                    default_value: None,
+                    fields: BTreeMap::new(),
+                    items: None,
+                    enum_values: Vec::new(),
+                    variants: Vec::new(),
+                },
+            )]),
+            items: None,
+            enum_values: Vec::new(),
+            variants: Vec::new(),
+        },
+        output_schema: Schema {
+            kind: SchemaKind::Object,
+            description: Some("output".to_owned()),
+            required: true,
+            nullable: false,
+            default_value: None,
+            fields: BTreeMap::from([(
+                "id".to_owned(),
+                Schema {
+                    kind: SchemaKind::String,
+                    description: None,
+                    required: true,
+                    nullable: false,
+                    default_value: None,
+                    fields: BTreeMap::new(),
+                    items: None,
+                    enum_values: Vec::new(),
+                    variants: Vec::new(),
+                },
+            )]),
+            items: None,
+            enum_values: Vec::new(),
+            variants: Vec::new(),
+        },
+        input_mapping: MappingSet {
+            rules: vec![MappingRule {
+                source: "$.mcp.email".to_owned(),
+                target: "$.request.body.email".to_owned(),
+                required: true,
+                default_value: None,
+                transform: None,
+                condition: None,
+                notes: None,
+            }],
+        },
+        output_mapping: MappingSet {
+            rules: vec![MappingRule {
+                source: "$.response.body.id".to_owned(),
+                target: "$.output.id".to_owned(),
+                required: true,
+                default_value: None,
+                transform: None,
+                condition: None,
+                notes: None,
+            }],
+        },
+        execution_config: ExecutionConfig {
+            timeout_ms: 10_000,
+            retry_policy: Some(RetryPolicy { max_attempts: 3 }),
+            auth_profile_ref: Some("auth_crank".into()),
+            headers: BTreeMap::from([("X-Request-Id".to_owned(), "static".to_owned())]),
+            ..ExecutionConfig::default()
+        },
+        tool_description: ToolDescription {
+            title: "Create lead".to_owned(),
+            description: "Creates CRM lead".to_owned(),
+            tags: vec!["crm".to_owned()],
+            examples: vec![ToolExample {
+                input: json!({ "email": "a@example.com" }),
+            }],
+        },
+        samples: Some(Samples {
+            input_json_sample_ref: Some("sample_input".into()),
+            output_json_sample_ref: Some("sample_output".into()),
+        }),
+        generated_draft: Some(GeneratedDraft {
+            status: GeneratedDraftStatus::Available,
+            source_types: vec!["input_json".to_owned(), "output_json".to_owned()],
+            generated_at: Some("2026-03-25T11:59:00Z".to_owned()),
+            input_schema_generated: true,
+            output_schema_generated: true,
+            input_mapping_generated: true,
+            output_mapping_generated: true,
+            warnings: Vec::new(),
+        }),
+        config_export: Some(ConfigExport {
+            format_version: "v1".to_owned(),
+            export_mode: ExportMode::Portable,
+        }),
+        wizard_state: Some(WizardState {
+            input_sample: Some(json!({ "email": format!("lead-{version}@example.com") })),
+            output_sample: Some(json!({ "id": format!("lead_{version}") })),
+            test_input: Some(json!({ "email": format!("test-v{version}@example.com") })),
+        }),
+        created_at: timestamp("2026-03-25T11:58:00Z"),
+        updated_at: timestamp(&format!("2026-03-25T12:{version:02}:00Z")),
+        published_at: None,
+    }
+}
+pub(super) fn test_agent(id: &str, status: AgentStatus) -> crank_core::Agent {
+    crank_core::Agent {
+        id: AgentId::new(id),
+        workspace_id: test_workspace_id(),
+        slug: format!("{id}_slug"),
+        display_name: format!("Display {id}"),
+        description: format!("Description {id}"),
+        status,
+        current_draft_version: 1,
+        latest_published_version: None,
+        created_at: timestamp("2026-03-25T11:58:00Z"),
+        updated_at: timestamp("2026-03-25T12:00:00Z"),
+        published_at: None,
+    }
+}
+
+pub(super) fn test_agent_version(
+    agent_id: &AgentId,
+    version: u32,
+    status: AgentStatus,
+) -> AgentVersion {
+    AgentVersion {
+        agent_id: agent_id.clone(),
+        version,
+        status,
+        instructions: json!({
+            "system": "triage tickets",
+            "guardrails": ["don't mutate state"]
+        }),
+        tool_selection_policy: json!({
+            "mode": "allow_list",
+            "max_tools": 8
+        }),
+        created_at: timestamp("2026-03-25T12:00:00Z"),
+    }
+}
+
+pub(super) fn test_invocation_log(
+    id: &str,
+    operation_id: &OperationId,
+    agent_id: Option<AgentId>,
+    status: crank_core::InvocationStatus,
+    duration_ms: u64,
+    created_at: &str,
+) -> InvocationLog {
+    InvocationLog {
+        id: crank_core::InvocationLogId::new(id),
+        workspace_id: test_workspace_id(),
+        agent_id,
+        operation_id: operation_id.clone(),
+        source: crank_core::InvocationSource::AgentToolCall,
+        level: crank_core::InvocationLevel::Info,
+        status,
+        tool_name: "create_lead".to_owned(),
+        message: "invocation".to_owned(),
+        request_id: Some(format!("req_{id}")),
+        status_code: Some(200),
+        duration_ms,
+        error_kind: None,
+        request_preview: json!({"input":"value"}),
+        response_preview: json!({"ok":true}),
+        created_at: timestamp(created_at),
+    }
+}
+
+pub(super) struct TestDatabase {
+    admin_pool: PgPool,
+    database_url: String,
+    schema: String,
+}
+
+impl TestDatabase {
+    pub(super) async fn new() -> Self {
+        let database_url = crank_test_support::postgres_database_url().await.to_owned();
+        let admin_pool = PgPoolOptions::new()
+            .max_connections(1)
+            .connect(&database_url)
+            .await
+            .unwrap();
+        let schema = crank_test_support::unique_schema_name("test_registry");
+
+        admin_pool
+            .execute(sqlx::query(&format!("create schema {schema}")))
+            .await
+            .unwrap();
+
+        Self {
+            admin_pool,
+            database_url,
+            schema,
+        }
+    }
+
+    pub(super) async fn registry(&self) -> PostgresRegistry {
+        PostgresRegistry::connect(&format!(
+            "{}?options=-csearch_path%3D{}",
+            self.database_url, self.schema
+        ))
+        .await
+        .unwrap()
+    }
+
+    pub(super) async fn cleanup(&self) {
+        self.admin_pool
+            .execute(sqlx::query(&format!(
+                "drop schema if exists {} cascade",
+                self.schema
+            )))
+            .await
+            .unwrap();
+    }
+}