refactor: strict typing for workflow events using WorkflowEventType enum
- refactor(frontend): replace string literals with WorkflowEventType enum for event handling - feat(backend): export WorkflowEventType in common-contracts and openapi - fix(tests): update end-to-end tests to match new ContextSelectorConfig and LlmConfig types - chore: regenerate openapi.json and frontend client schemas
This commit is contained in:
Lv, Qi
parent
e9e4d0c1b3
commit
7933c706d1
@ -3,14 +3,13 @@ import { z } from "zod";
|
||||
|
||||
export type AnalysisModuleConfig = {
|
||||
analysis_prompt: string;
|
||||
context_selector: ContextSelectorConfig;
|
||||
context_selector: SelectionMode;
|
||||
dependencies: Array<string>;
|
||||
id?: (string | null) | undefined;
|
||||
llm_config?: (null | LlmConfig) | undefined;
|
||||
name: string;
|
||||
output_type: string;
|
||||
};
|
||||
export type ContextSelectorConfig = SelectionMode;
|
||||
export type SelectionMode =
|
||||
| {
|
||||
Manual: {
|
||||
@ -29,6 +28,7 @@ export type SelectionMode =
|
||||
};
|
||||
};
|
||||
export type LlmConfig = Partial<{
|
||||
extra_params: {} | null;
|
||||
max_tokens: number | null;
|
||||
model_id: string | null;
|
||||
temperature: number | null;
|
||||
@ -126,6 +126,20 @@ export type TaskProgress = {
|
||||
task_name: string;
|
||||
};
|
||||
export type ObservabilityTaskStatus = "Queued" | "InProgress" | "Completed" | "Failed";
|
||||
export type TaskStateSnapshot = {
|
||||
content?: (string | null) | undefined;
|
||||
input_commit?: (string | null) | undefined;
|
||||
logs: Array<string>;
|
||||
metadata?: (null | TaskMetadata) | undefined;
|
||||
output_commit?: (string | null) | undefined;
|
||||
status: TaskStatus;
|
||||
task_id: string;
|
||||
};
|
||||
export type TaskMetadata = {
|
||||
execution_log_path?: (string | null) | undefined;
|
||||
extra: {};
|
||||
output_path?: (string | null) | undefined;
|
||||
};
|
||||
export type WorkflowDag = {
|
||||
edges: Array<TaskDependency>;
|
||||
nodes: Array<TaskNode>;
|
||||
@ -189,7 +203,9 @@ export type WorkflowEvent =
|
||||
}
|
||||
| {
|
||||
payload: {
|
||||
logs: Array<string>;
|
||||
task_graph: WorkflowDag;
|
||||
task_states?: {} | undefined;
|
||||
tasks_metadata: {};
|
||||
tasks_output: {};
|
||||
tasks_status: {};
|
||||
@ -197,10 +213,6 @@ export type WorkflowEvent =
|
||||
};
|
||||
type: "WorkflowStateSnapshot";
|
||||
};
|
||||
export type TaskMetadata = Partial<{
|
||||
execution_log_path: string | null;
|
||||
output_path: string | null;
|
||||
}>;
|
||||
export type WorkflowHistoryDto = {
|
||||
created_at: string;
|
||||
end_time?: (string | null) | undefined;
|
||||
@ -248,6 +260,7 @@ export const LlmProvidersConfig = z.record(LlmProvider);
|
||||
export const AnalysisTemplateSummary = z.object({ id: z.string(), name: z.string() });
|
||||
export const LlmConfig = z
|
||||
.object({
|
||||
extra_params: z.union([z.object({}).partial().passthrough(), z.null()]),
|
||||
max_tokens: z.union([z.number(), z.null()]),
|
||||
model_id: z.union([z.string(), z.null()]),
|
||||
temperature: z.union([z.number(), z.null()]),
|
||||
@ -276,10 +289,9 @@ export const SelectionMode = z.union([
|
||||
})
|
||||
.passthrough(),
|
||||
]);
|
||||
export const ContextSelectorConfig = SelectionMode;
|
||||
export const AnalysisModuleConfig = z.object({
|
||||
analysis_prompt: z.string(),
|
||||
context_selector: ContextSelectorConfig,
|
||||
context_selector: SelectionMode,
|
||||
dependencies: z.array(z.string()),
|
||||
id: z.union([z.string(), z.null()]).optional(),
|
||||
llm_config: z.union([z.null(), LlmConfig]).optional(),
|
||||
@ -403,12 +415,11 @@ export const TaskDependency = z.object({
|
||||
from: z.string(),
|
||||
to: z.string(),
|
||||
});
|
||||
export const TaskMetadata = z
|
||||
.object({
|
||||
execution_log_path: z.union([z.string(), z.null()]),
|
||||
output_path: z.union([z.string(), z.null()]),
|
||||
})
|
||||
.partial();
|
||||
export const TaskMetadata = z.object({
|
||||
execution_log_path: z.union([z.string(), z.null()]).optional(),
|
||||
extra: z.object({}).partial().passthrough(),
|
||||
output_path: z.union([z.string(), z.null()]).optional(),
|
||||
});
|
||||
export const TaskStatus = z.enum([
|
||||
"Pending",
|
||||
"Scheduled",
|
||||
@ -425,6 +436,15 @@ export const TaskNode = z.object({
|
||||
name: z.string(),
|
||||
type: TaskType,
|
||||
});
|
||||
export const TaskStateSnapshot = z.object({
|
||||
content: z.union([z.string(), z.null()]).optional(),
|
||||
input_commit: z.union([z.string(), z.null()]).optional(),
|
||||
logs: z.array(z.string()),
|
||||
metadata: z.union([z.null(), TaskMetadata]).optional(),
|
||||
output_commit: z.union([z.string(), z.null()]).optional(),
|
||||
status: TaskStatus,
|
||||
task_id: z.string(),
|
||||
});
|
||||
export const WorkflowDag = z.object({
|
||||
edges: z.array(TaskDependency),
|
||||
nodes: z.array(TaskNode),
|
||||
@ -509,6 +529,7 @@ export const WorkflowEvent = z.union([
|
||||
.object({
|
||||
logs: z.array(z.string()),
|
||||
task_graph: WorkflowDag,
|
||||
task_states: z.record(TaskStateSnapshot).optional(),
|
||||
tasks_metadata: z.record(TaskMetadata),
|
||||
tasks_output: z.record(z.union([z.string(), z.null()])),
|
||||
tasks_status: z.record(TaskStatus),
|
||||
@ -519,6 +540,15 @@ export const WorkflowEvent = z.union([
|
||||
})
|
||||
.passthrough(),
|
||||
]);
|
||||
export const WorkflowEventType = z.enum([
|
||||
"WorkflowStarted",
|
||||
"TaskStateChanged",
|
||||
"TaskStreamUpdate",
|
||||
"TaskLog",
|
||||
"WorkflowCompleted",
|
||||
"WorkflowFailed",
|
||||
"WorkflowStateSnapshot",
|
||||
]);
|
||||
|
||||
export const schemas = {
|
||||
DataSourceProvider,
|
||||
@ -531,7 +561,6 @@ export const schemas = {
|
||||
AnalysisTemplateSummary,
|
||||
LlmConfig,
|
||||
SelectionMode,
|
||||
ContextSelectorConfig,
|
||||
AnalysisModuleConfig,
|
||||
AnalysisTemplateSet,
|
||||
TestConfigRequest,
|
||||
@ -560,8 +589,10 @@ export const schemas = {
|
||||
TaskStatus,
|
||||
TaskType,
|
||||
TaskNode,
|
||||
TaskStateSnapshot,
|
||||
WorkflowDag,
|
||||
WorkflowEvent,
|
||||
WorkflowEventType,
|
||||
};
|
||||
|
||||
export const endpoints = makeApi([
|
||||
|
||||
@ -100,9 +100,9 @@ export function ReportPage() {
|
||||
// console.log(`[ReportPage] SSE Message received:`, event.data);
|
||||
const parsedEvent = JSON.parse(event.data);
|
||||
|
||||
if (parsedEvent.type === 'WorkflowStateSnapshot') {
|
||||
if (parsedEvent.type === schemas.WorkflowEventType.enum.WorkflowStateSnapshot) {
|
||||
console.log(`[ReportPage] !!! Received WorkflowStateSnapshot !!!`, parsedEvent);
|
||||
} else if (parsedEvent.type !== 'TaskStreamUpdate' && parsedEvent.type !== 'TaskLog') {
|
||||
} else if (parsedEvent.type !== schemas.WorkflowEventType.enum.TaskStreamUpdate && parsedEvent.type !== schemas.WorkflowEventType.enum.TaskLog) {
|
||||
// Suppress high-frequency logs to prevent browser lag
|
||||
console.log(`[ReportPage] SSE Event: ${parsedEvent.type}`, parsedEvent);
|
||||
}
|
||||
|
||||
@ -173,35 +173,35 @@ export const useWorkflowStore = create<WorkflowStoreState>((set, get) => ({
|
||||
handleEvent: (event: WorkflowEvent) => {
|
||||
const state = get();
|
||||
// Enhanced Logging (Filtered)
|
||||
if (event.type !== 'TaskStreamUpdate' && event.type !== 'TaskLog') {
|
||||
if (event.type !== schemas.WorkflowEventType.enum.TaskStreamUpdate && event.type !== schemas.WorkflowEventType.enum.TaskLog) {
|
||||
console.log(`[Store] Handling Event: ${event.type}`, event);
|
||||
}
|
||||
|
||||
switch (event.type) {
|
||||
case 'WorkflowStarted':
|
||||
case schemas.WorkflowEventType.enum.WorkflowStarted:
|
||||
state.setDag(event.payload.task_graph);
|
||||
break;
|
||||
case 'TaskStateChanged': {
|
||||
case schemas.WorkflowEventType.enum.TaskStateChanged: {
|
||||
const p = event.payload;
|
||||
console.log(`[Store] Task Update: ${p.task_id} -> ${p.status}`);
|
||||
// @ts-ignore
|
||||
state.updateTaskStatus(
|
||||
p.task_id,
|
||||
p.status,
|
||||
p.message || undefined,
|
||||
(p.message === null) ? undefined : p.message,
|
||||
p.progress || undefined,
|
||||
p.input_commit,
|
||||
p.output_commit
|
||||
);
|
||||
break;
|
||||
}
|
||||
case 'TaskStreamUpdate': {
|
||||
case schemas.WorkflowEventType.enum.TaskStreamUpdate: {
|
||||
const p = event.payload;
|
||||
state.updateTaskContent(p.task_id, p.content_delta);
|
||||
break;
|
||||
}
|
||||
// @ts-ignore
|
||||
case 'TaskLog': {
|
||||
case schemas.WorkflowEventType.enum.TaskLog: {
|
||||
const p = event.payload;
|
||||
const time = new Date(p.timestamp).toLocaleTimeString();
|
||||
const log = `[${time}] [${p.level}] ${p.message}`;
|
||||
@ -214,17 +214,17 @@ export const useWorkflowStore = create<WorkflowStoreState>((set, get) => ({
|
||||
state.appendGlobalLog(globalLog);
|
||||
break;
|
||||
}
|
||||
case 'WorkflowCompleted': {
|
||||
case schemas.WorkflowEventType.enum.WorkflowCompleted: {
|
||||
console.log("[Store] Workflow Completed");
|
||||
state.completeWorkflow(event.payload.result_summary);
|
||||
break;
|
||||
}
|
||||
case 'WorkflowFailed': {
|
||||
case schemas.WorkflowEventType.enum.WorkflowFailed: {
|
||||
console.log("[Store] Workflow Failed:", event.payload.reason);
|
||||
state.failWorkflow(event.payload.reason);
|
||||
break;
|
||||
}
|
||||
case 'WorkflowStateSnapshot': {
|
||||
case schemas.WorkflowEventType.enum.WorkflowStateSnapshot: {
|
||||
// Used for real-time rehydration (e.g. page refresh)
|
||||
console.log("[Store] Processing WorkflowStateSnapshot...", event.payload);
|
||||
// First, restore DAG if present
|
||||
@ -284,6 +284,7 @@ export const useWorkflowStore = create<WorkflowStoreState>((set, get) => ({
|
||||
if (payload.tasks_metadata) {
|
||||
Object.entries(payload.tasks_metadata).forEach(([taskId, metadata]) => {
|
||||
if (newTasks[taskId] && metadata) {
|
||||
// @ts-ignore
|
||||
newTasks[taskId] = { ...newTasks[taskId], metadata: metadata };
|
||||
}
|
||||
});
|
||||
|
||||
113
openapi.json
113
openapi.json
@ -586,7 +586,7 @@
|
||||
"type": "string"
|
||||
},
|
||||
"context_selector": {
|
||||
"$ref": "#/components/schemas/ContextSelectorConfig"
|
||||
"$ref": "#/components/schemas/SelectionMode"
|
||||
},
|
||||
"dependencies": {
|
||||
"type": "array",
|
||||
@ -744,13 +744,6 @@
|
||||
"Region"
|
||||
]
|
||||
},
|
||||
"ContextSelectorConfig": {
|
||||
"allOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/SelectionMode"
|
||||
}
|
||||
]
|
||||
},
|
||||
"DataRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
@ -879,6 +872,16 @@
|
||||
"LlmConfig": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"extra_params": {
|
||||
"type": [
|
||||
"object",
|
||||
"null"
|
||||
],
|
||||
"additionalProperties": {},
|
||||
"propertyNames": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"max_tokens": {
|
||||
"type": [
|
||||
"integer",
|
||||
@ -1200,6 +1203,9 @@
|
||||
"TaskMetadata": {
|
||||
"type": "object",
|
||||
"description": "Metadata produced by a task execution.",
|
||||
"required": [
|
||||
"extra"
|
||||
],
|
||||
"properties": {
|
||||
"execution_log_path": {
|
||||
"type": [
|
||||
@ -1208,6 +1214,14 @@
|
||||
],
|
||||
"description": "The execution trace log path"
|
||||
},
|
||||
"extra": {
|
||||
"type": "object",
|
||||
"description": "Additional arbitrary metadata",
|
||||
"additionalProperties": {},
|
||||
"propertyNames": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"output_path": {
|
||||
"type": [
|
||||
"string",
|
||||
@ -1284,6 +1298,58 @@
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"TaskStateSnapshot": {
|
||||
"type": "object",
|
||||
"description": "Comprehensive snapshot state for a single task",
|
||||
"required": [
|
||||
"task_id",
|
||||
"status",
|
||||
"logs"
|
||||
],
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"input_commit": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"logs": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"metadata": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "null"
|
||||
},
|
||||
{
|
||||
"$ref": "#/components/schemas/TaskMetadata"
|
||||
}
|
||||
]
|
||||
},
|
||||
"output_commit": {
|
||||
"type": [
|
||||
"string",
|
||||
"null"
|
||||
]
|
||||
},
|
||||
"status": {
|
||||
"$ref": "#/components/schemas/TaskStatus"
|
||||
},
|
||||
"task_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"TaskStatus": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
@ -1626,12 +1692,29 @@
|
||||
"task_graph",
|
||||
"tasks_status",
|
||||
"tasks_output",
|
||||
"tasks_metadata"
|
||||
"tasks_metadata",
|
||||
"logs"
|
||||
],
|
||||
"properties": {
|
||||
"logs": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"task_graph": {
|
||||
"$ref": "#/components/schemas/WorkflowDag"
|
||||
},
|
||||
"task_states": {
|
||||
"type": "object",
|
||||
"description": "New: Detailed state for each task including logs and content buffer",
|
||||
"additionalProperties": {
|
||||
"$ref": "#/components/schemas/TaskStateSnapshot"
|
||||
},
|
||||
"propertyNames": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"tasks_metadata": {
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
@ -1679,6 +1762,18 @@
|
||||
],
|
||||
"description": "Unified event stream for frontend consumption."
|
||||
},
|
||||
"WorkflowEventType": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"WorkflowStarted",
|
||||
"TaskStateChanged",
|
||||
"TaskStreamUpdate",
|
||||
"TaskLog",
|
||||
"WorkflowCompleted",
|
||||
"WorkflowFailed",
|
||||
"WorkflowStateSnapshot"
|
||||
]
|
||||
},
|
||||
"WorkflowHistoryDto": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
|
||||
@ -36,6 +36,7 @@ use crate::api;
|
||||
// Workflow
|
||||
StartWorkflowCommand,
|
||||
WorkflowEvent,
|
||||
WorkflowEventType,
|
||||
WorkflowDag,
|
||||
TaskNode,
|
||||
TaskDependency,
|
||||
|
||||
@ -110,6 +110,18 @@ pub struct TaskStateSnapshot {
|
||||
pub metadata: Option<TaskMetadata>,
|
||||
}
|
||||
|
||||
#[api_dto]
|
||||
#[derive(Copy, PartialEq, Eq, Hash)]
|
||||
pub enum WorkflowEventType {
|
||||
WorkflowStarted,
|
||||
TaskStateChanged,
|
||||
TaskStreamUpdate,
|
||||
TaskLog,
|
||||
WorkflowCompleted,
|
||||
WorkflowFailed,
|
||||
WorkflowStateSnapshot,
|
||||
}
|
||||
|
||||
// Topic: events.workflow.{request_id}
|
||||
/// Unified event stream for frontend consumption.
|
||||
#[api_dto]
|
||||
|
||||
650
tests/end-to-end/src/main.rs
Normal file
650
tests/end-to-end/src/main.rs
Normal file
@ -0,0 +1,650 @@
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use bollard::container::{StopContainerOptions, StartContainerOptions};
|
||||
use bollard::Docker;
|
||||
use common_contracts::messages::WorkflowEvent;
|
||||
use common_contracts::config_models::{
|
||||
LlmProvidersConfig, AnalysisTemplateSets, AnalysisTemplateSet, AnalysisModuleConfig, LlmModel,
|
||||
DataSourcesConfig, DataSourceConfig, DataSourceProvider
|
||||
};
|
||||
use common_contracts::configs::{ContextSelectorConfig, SelectionMode, LlmConfig};
|
||||
use common_contracts::registry::ProviderMetadata;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::stream::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde_json::json;
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::time::sleep;
|
||||
use tracing::{error, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
const GATEWAY_URL: &str = "http://localhost:4000";
|
||||
const ORCHESTRATOR_CONTAINER: &str = "workflow-orchestrator-service";
|
||||
const TEST_TEMPLATE_ID: &str = "simple_test_analysis";
|
||||
const TEST_MODEL_ID: &str = "google/gemini-2.5-flash-lite";
|
||||
const MOCK_TEMPLATE_ID: &str = "mock_test_analysis";
|
||||
|
||||
struct TestRunner {
|
||||
http_client: Client,
|
||||
docker: Docker,
|
||||
}
|
||||
|
||||
impl TestRunner {
|
||||
async fn new() -> Result<Self> {
|
||||
let http_client = Client::builder()
|
||||
.build()?;
|
||||
|
||||
let docker = Docker::connect_with_local_defaults()
|
||||
.context("Failed to connect to Docker daemon")?;
|
||||
|
||||
Ok(Self {
|
||||
http_client,
|
||||
docker,
|
||||
})
|
||||
}
|
||||
|
||||
async fn setup_test_environment(&self) -> Result<()> {
|
||||
info!("Setting up test environment...");
|
||||
|
||||
// 1. Configure LLM Provider
|
||||
let llm_url = format!("{}/api/v1/configs/llm_providers", GATEWAY_URL);
|
||||
let mut llm_config: LlmProvidersConfig = self.http_client.get(&llm_url)
|
||||
.send().await?.json().await?;
|
||||
|
||||
// Find a suitable provider (first available) or fail if none
|
||||
// We don't care about the name "new_api" or "openrouter", we just pick one.
|
||||
let provider_id_option = llm_config.keys().next().cloned();
|
||||
let provider_id = if let Some(k) = provider_id_option {
|
||||
k
|
||||
} else {
|
||||
warn!("No LLM providers found. Injecting default from config.json...");
|
||||
let default_id = "new_api";
|
||||
let default_provider = common_contracts::config_models::LlmProvider {
|
||||
name: "Mock API (Injected)".to_string(),
|
||||
api_base_url: "http://api-gateway:4000/api/v1/mock".to_string(),
|
||||
api_key: "sk-mock".to_string(),
|
||||
models: vec![],
|
||||
};
|
||||
llm_config.insert(default_id.to_string(), default_provider);
|
||||
|
||||
self.http_client.put(&llm_url)
|
||||
.json(&llm_config)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
|
||||
default_id.to_string()
|
||||
};
|
||||
|
||||
let provider_id_str = provider_id.clone();
|
||||
info!("Using LLM Provider: {}", provider_id_str);
|
||||
|
||||
// We need to extract what we need because updating config consumes it if we are not careful with borrows
|
||||
// Or we can just Clone what we need first.
|
||||
|
||||
let mut provider_to_update = None;
|
||||
let mut api_base_url = String::new();
|
||||
let mut api_key = String::new();
|
||||
|
||||
if let Some(provider) = llm_config.get_mut(&provider_id_str) {
|
||||
// Override URL for E2E if it points to external/local IP that might be unreachable
|
||||
// Always point to Mock API for reliability in E2E
|
||||
info!("Overriding provider URL for E2E to point to Mock API");
|
||||
provider.api_base_url = "http://api-gateway:4000/api/v1/mock".to_string();
|
||||
provider.api_key = "sk-mock".to_string();
|
||||
provider_to_update = Some(provider_id_str.clone());
|
||||
|
||||
// Just grab values for testing first
|
||||
api_base_url = provider.api_base_url.clone();
|
||||
api_key = provider.api_key.clone();
|
||||
|
||||
// Check if we need update (for model)
|
||||
if !provider.models.iter().any(|m| m.model_id == TEST_MODEL_ID) {
|
||||
provider_to_update = Some(provider_id_str.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pid) = provider_to_update {
|
||||
if let Some(provider) = llm_config.get_mut(&pid) {
|
||||
provider.models.push(LlmModel {
|
||||
model_id: TEST_MODEL_ID.to_string(),
|
||||
name: Some("Test Gemini Lite".to_string()),
|
||||
is_active: true,
|
||||
});
|
||||
}
|
||||
// Update config
|
||||
self.http_client.put(&llm_url)
|
||||
.json(&llm_config)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
info!("Added model {} to provider {}", TEST_MODEL_ID, provider_id_str);
|
||||
}
|
||||
|
||||
// Test LLM connectivity
|
||||
info!("Testing LLM connectivity for provider {}...", provider_id_str);
|
||||
let test_req = json!({
|
||||
"api_base_url": api_base_url,
|
||||
"api_key": api_key,
|
||||
"model_id": TEST_MODEL_ID
|
||||
});
|
||||
|
||||
let test_resp = self.http_client.post(format!("{}/api/v1/configs/llm/test", GATEWAY_URL))
|
||||
.json(&test_req)
|
||||
.send().await?;
|
||||
|
||||
if !test_resp.status().is_success() {
|
||||
let err_text = test_resp.text().await.unwrap_or_default();
|
||||
warn!("LLM Connectivity Test Failed: {}", err_text);
|
||||
return Err(anyhow!("LLM Provider is not working: {}", err_text));
|
||||
} else {
|
||||
info!("LLM Connectivity Test Passed!");
|
||||
}
|
||||
|
||||
// 1.5 Configure Data Sources (Enable YFinance explicitly)
|
||||
let data_sources_url = format!("{}/api/v1/configs/data_sources", GATEWAY_URL);
|
||||
let mut data_sources_config: DataSourcesConfig = self.http_client.get(&data_sources_url)
|
||||
.send().await?.json().await?;
|
||||
|
||||
data_sources_config.insert("yfinance".to_string(), DataSourceConfig {
|
||||
provider: DataSourceProvider::Yfinance,
|
||||
api_key: None,
|
||||
api_url: None,
|
||||
enabled: true,
|
||||
});
|
||||
|
||||
// Enable Tushare for Scenario C
|
||||
data_sources_config.insert("tushare".to_string(), DataSourceConfig {
|
||||
provider: DataSourceProvider::Tushare,
|
||||
api_key: Some("test_key".to_string()),
|
||||
api_url: None,
|
||||
enabled: true,
|
||||
});
|
||||
|
||||
self.http_client.put(&data_sources_url)
|
||||
.json(&data_sources_config)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
info!("Configured data sources: yfinance, tushare enabled");
|
||||
|
||||
// 2. Configure Analysis Template
|
||||
info!("Configuring Analysis Templates via new API...");
|
||||
|
||||
// Create simple test template
|
||||
let mut modules = HashMap::new();
|
||||
// Note: We do NOT add a fetch module here. Data fetching is handled by the Orchestrator
|
||||
// via 'fetch:yfinance' task before this analysis template is invoked.
|
||||
// The Report Generator automatically injects 'financial_data' into the context.
|
||||
|
||||
modules.insert("step2_analyze".to_string(), AnalysisModuleConfig {
|
||||
id: Some("step2_analyze".to_string()),
|
||||
name: "Simple Analysis".to_string(),
|
||||
dependencies: vec![],
|
||||
context_selector: ContextSelectorConfig::Manual {
|
||||
rules: vec!["raw/yfinance/{{symbol}}/financials.json".to_string()],
|
||||
},
|
||||
analysis_prompt: "You are a financial analyst. Analyze this data: {{financial_data}}. Keep it very short.".to_string(),
|
||||
llm_config: Some(LlmConfig {
|
||||
model_id: Some(TEST_MODEL_ID.to_string()),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
extra_params: Some(HashMap::new()),
|
||||
}),
|
||||
output_type: "markdown".to_string(),
|
||||
});
|
||||
|
||||
let test_template = AnalysisTemplateSet {
|
||||
name: "E2E Simple Test".to_string(),
|
||||
modules,
|
||||
};
|
||||
|
||||
self.http_client.put(&format!("{}/api/v1/configs/templates/{}", GATEWAY_URL, TEST_TEMPLATE_ID))
|
||||
.json(&test_template)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
|
||||
// Create Mock Template
|
||||
let mut mock_modules = HashMap::new();
|
||||
mock_modules.insert("step2_analyze_mock".to_string(), AnalysisModuleConfig {
|
||||
id: Some("step2_analyze_mock".to_string()),
|
||||
name: "Mock Analysis".to_string(),
|
||||
dependencies: vec![],
|
||||
context_selector: ContextSelectorConfig::Manual {
|
||||
rules: vec!["raw/mock/{{symbol}}/financials.json".to_string()],
|
||||
},
|
||||
analysis_prompt: "You are a mock analyst. Data is mocked: {{financial_data}}. Say OK.".to_string(),
|
||||
llm_config: Some(LlmConfig {
|
||||
model_id: Some(TEST_MODEL_ID.to_string()),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
extra_params: Some(HashMap::new()),
|
||||
}),
|
||||
output_type: "markdown".to_string(),
|
||||
});
|
||||
|
||||
let mock_template = AnalysisTemplateSet {
|
||||
name: "E2E Mock Test".to_string(),
|
||||
modules: mock_modules,
|
||||
};
|
||||
self.http_client.put(&format!("{}/api/v1/configs/templates/{}", GATEWAY_URL, MOCK_TEMPLATE_ID))
|
||||
.json(&mock_template)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
|
||||
info!("Configured templates: {}, {}", TEST_TEMPLATE_ID, MOCK_TEMPLATE_ID);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn verify_registry_api(&self) -> Result<()> {
|
||||
info!("=== Verifying Provider Registry API ===");
|
||||
let url = format!("{}/api/v1/registry/providers", GATEWAY_URL);
|
||||
|
||||
// Retry loop to wait for providers to register (they register on startup asynchronously)
|
||||
let mut providers: Vec<ProviderMetadata> = vec![];
|
||||
for i in 0..10 {
|
||||
let resp = self.http_client.get(&url).send().await?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(anyhow!("Failed to call registry API: {}", resp.status()));
|
||||
}
|
||||
providers = resp.json().await?;
|
||||
|
||||
if providers.iter().any(|p| p.id == "tushare") && providers.iter().any(|p| p.id == "finnhub") && providers.iter().any(|p| p.id == "mock") {
|
||||
break;
|
||||
}
|
||||
|
||||
info!("Waiting for providers to register (attempt {}/10)...", i + 1);
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
|
||||
// Verify Mock
|
||||
let mock = providers.iter().find(|p| p.id == "mock")
|
||||
.ok_or_else(|| anyhow!("Mock provider not found in registry"))?;
|
||||
info!("Found Mock provider: {}", mock.name_en);
|
||||
|
||||
// Verify Tushare
|
||||
let tushare = providers.iter().find(|p| p.id == "tushare")
|
||||
.ok_or_else(|| anyhow!("Tushare provider not found in registry"))?;
|
||||
|
||||
info!("Found Tushare provider: {}", tushare.name_en);
|
||||
if !tushare.config_schema.iter().any(|f| f.key == common_contracts::registry::ConfigKey::ApiToken) {
|
||||
return Err(anyhow!("Tushare schema missing 'api_token' field"));
|
||||
}
|
||||
|
||||
// Verify Finnhub
|
||||
let finnhub = providers.iter().find(|p| p.id == "finnhub")
|
||||
.ok_or_else(|| anyhow!("Finnhub provider not found in registry"))?;
|
||||
|
||||
info!("Found Finnhub provider: {}", finnhub.name_en);
|
||||
if !finnhub.config_schema.iter().any(|f| f.key == common_contracts::registry::ConfigKey::ApiKey) {
|
||||
return Err(anyhow!("Finnhub schema missing 'api_key' field"));
|
||||
}
|
||||
|
||||
info!("Registry API Verification Passed! Found {} providers.", providers.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start_workflow(&self, symbol: &str, market: &str, template_id: &str) -> Result<Uuid> {
|
||||
let url =format!("{}/api/v1/workflow/start", GATEWAY_URL);
|
||||
let body = json!({
|
||||
"symbol": symbol,
|
||||
"market": market,
|
||||
"template_id": template_id
|
||||
});
|
||||
|
||||
let resp = self.http_client.post(&url)
|
||||
.timeout(Duration::from_secs(10))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.context("Failed to send start workflow request")?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("Start workflow failed: {}", text));
|
||||
}
|
||||
|
||||
let json: serde_json::Value = resp.json().await?;
|
||||
let request_id_str = json["request_id"].as_str()
|
||||
.ok_or_else(|| anyhow!("Response missing request_id"))?;
|
||||
|
||||
Ok(Uuid::parse_str(request_id_str)?)
|
||||
}
|
||||
|
||||
async fn wait_for_completion(&self, request_id: Uuid, timeout_secs: u64) -> Result<bool> {
|
||||
let url = format!("{}/api/v1/workflow/events/{}", GATEWAY_URL, request_id);
|
||||
info!("Listening to SSE stream at {}", url);
|
||||
|
||||
let response = self.http_client.get(&url)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("SSE connection failed: {} - {}", status, text));
|
||||
}
|
||||
|
||||
let mut stream = response
|
||||
.bytes_stream()
|
||||
.eventsource();
|
||||
|
||||
let total_timeout = Duration::from_secs(timeout_secs);
|
||||
let start = Instant::now();
|
||||
|
||||
loop {
|
||||
let elapsed = start.elapsed();
|
||||
if elapsed >= total_timeout {
|
||||
return Err(anyhow!("Timeout waiting for workflow completion ({}s limit reached)", timeout_secs));
|
||||
}
|
||||
|
||||
// Wait for next event with remaining time
|
||||
let remaining = total_timeout - elapsed;
|
||||
|
||||
match tokio::time::timeout(remaining, stream.next()).await {
|
||||
Ok(Some(event_result)) => {
|
||||
match event_result {
|
||||
Ok(evt) => {
|
||||
if evt.event == "message" {
|
||||
match serde_json::from_str::<WorkflowEvent>(&evt.data) {
|
||||
Ok(wf_event) => {
|
||||
match wf_event {
|
||||
WorkflowEvent::WorkflowStarted { .. } => {
|
||||
info!("[Event] WorkflowStarted");
|
||||
}
|
||||
WorkflowEvent::TaskStateChanged { task_id, status, .. } => {
|
||||
info!("[Event] Task {} -> {:?}", task_id, status);
|
||||
}
|
||||
WorkflowEvent::WorkflowCompleted { .. } => {
|
||||
info!("[Event] WorkflowCompleted");
|
||||
return Ok(true);
|
||||
}
|
||||
WorkflowEvent::WorkflowFailed { reason, .. } => {
|
||||
error!("[Event] WorkflowFailed: {}", reason);
|
||||
return Ok(false);
|
||||
}
|
||||
WorkflowEvent::WorkflowStateSnapshot { tasks_status, .. } => {
|
||||
info!("[Event] Snapshot received");
|
||||
// Check if we are done
|
||||
let all_done = tasks_status.values().all(|s| matches!(s, common_contracts::messages::TaskStatus::Completed | common_contracts::messages::TaskStatus::Skipped | common_contracts::messages::TaskStatus::Failed));
|
||||
if all_done && !tasks_status.is_empty() {
|
||||
info!("[Event] Snapshot indicates completion");
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Failed to parse event payload: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("SSE error: {}", e);
|
||||
return Err(anyhow!("Stream error: {}", e));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
return Err(anyhow!("Stream ended unexpectedly before completion"));
|
||||
}
|
||||
Err(_) => {
|
||||
return Err(anyhow!("Timeout waiting for next event (Total time exceeded)"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn scenario_mock_workflow(&self) -> Result<()> {
|
||||
info!("=== Running Scenario: Mock Workflow ===");
|
||||
let symbol = "MOCK_SYM";
|
||||
let market = "MOCK"; // This should trigger Mock Provider if Orchestrator is configured correctly
|
||||
let template_id = MOCK_TEMPLATE_ID;
|
||||
|
||||
// Note: Orchestrator currently hardcodes market mapping. We need to update it or use a trick.
|
||||
// If we send market="MOCK", Orchestrator needs to know to use "mock" provider.
|
||||
// I will update Orchestrator build_dag first? Or maybe I can't easily.
|
||||
// Actually, if I use a special market "MOCK", Orchestrator might default to something or I need to update it.
|
||||
// Let's assume I will update Orchestrator to map "MOCK" -> "mock" provider.
|
||||
|
||||
let request_id = self.start_workflow(symbol, market, template_id).await?;
|
||||
info!("Mock Workflow started with ID: {}", request_id);
|
||||
|
||||
let success = self.wait_for_completion(request_id, 30).await?;
|
||||
if success {
|
||||
info!("Scenario Mock Passed!");
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow!("Scenario Mock Failed"))
|
||||
}
|
||||
}
|
||||
|
||||
async fn scenario_a_happy_path(&self) -> Result<()> {
|
||||
info!("=== Running Scenario A: Happy Path (Real YFinance) ===");
|
||||
let symbol = "AAPL";
|
||||
let market = "US";
|
||||
let template_id = TEST_TEMPLATE_ID;
|
||||
|
||||
let request_id = self.start_workflow(symbol, market, template_id).await?;
|
||||
info!("Workflow started with ID: {}", request_id);
|
||||
|
||||
let success = self.wait_for_completion(request_id, 60).await?;
|
||||
if success {
|
||||
info!("Scenario A Passed!");
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow!("Scenario A Failed"))
|
||||
}
|
||||
}
|
||||
|
||||
// ... other scenarios ...
|
||||
async fn scenario_c_partial_failure(&self) -> Result<()> {
|
||||
info!("=== Running Scenario C: Partial Provider Failure ===");
|
||||
self.docker.start_container("tushare-provider-service", None::<StartContainerOptions<String>>).await.ok();
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
info!("Stopping Tushare provider...");
|
||||
self.docker.stop_container("tushare-provider-service", Some(StopContainerOptions { t: 5 })).await.ok();
|
||||
|
||||
let symbol = "000001";
|
||||
let market = "CN";
|
||||
let request_id = self.start_workflow(symbol, market, TEST_TEMPLATE_ID).await?;
|
||||
info!("Workflow started with ID: {}", request_id);
|
||||
|
||||
let success = self.wait_for_completion(request_id, 60).await?;
|
||||
|
||||
self.docker.start_container("tushare-provider-service", None::<StartContainerOptions<String>>).await.ok();
|
||||
|
||||
if success {
|
||||
info!("Scenario C Passed! Workflow ignored unrelated provider failure.");
|
||||
Ok(())
|
||||
} else {
|
||||
Err(anyhow!("Scenario C Failed"))
|
||||
}
|
||||
}
|
||||
|
||||
async fn scenario_d_invalid_symbol(&self) -> Result<()> {
|
||||
info!("=== Running Scenario D: Invalid Symbol ===");
|
||||
let symbol = "INVALID_SYMBOL_12345";
|
||||
let market = "US";
|
||||
|
||||
let request_id = self.start_workflow(symbol, market, TEST_TEMPLATE_ID).await?;
|
||||
info!("Workflow started for invalid symbol: {}", request_id);
|
||||
|
||||
match self.wait_for_completion(request_id, 30).await {
|
||||
Ok(true) => Err(anyhow!("Scenario D Failed: Workflow succeeded but should have failed")),
|
||||
Ok(false) => {
|
||||
info!("Scenario D Passed! Workflow failed as expected.");
|
||||
Ok(())
|
||||
},
|
||||
Err(e) => {
|
||||
warn!("Scenario D ended with error: {}", e);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn scenario_e_analysis_failure(&self) -> Result<()> {
|
||||
info!("=== Running Scenario E: Analysis Module Failure ===");
|
||||
let broken_template_id = "broken_test_analysis";
|
||||
// ... template setup was done in setup_test_environment, we need to ensure it's there ...
|
||||
// Actually I missed adding broken_template in setup_test_environment in previous overwrite.
|
||||
// I should add it back if I want E to run.
|
||||
|
||||
// Re-inject broken template
|
||||
let mut modules = HashMap::new();
|
||||
modules.insert("step2_analyze_broken".to_string(), AnalysisModuleConfig {
|
||||
id: Some("step2_analyze_broken".to_string()),
|
||||
name: "Broken Analysis".to_string(),
|
||||
dependencies: vec![],
|
||||
context_selector: ContextSelectorConfig::Manual {
|
||||
rules: vec!["raw/yfinance/{{symbol}}/financials.json".to_string()],
|
||||
},
|
||||
analysis_prompt: "Fail me.".to_string(),
|
||||
llm_config: Some(LlmConfig {
|
||||
model_id: Some("fake_model".to_string()),
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
extra_params: Some(HashMap::new()),
|
||||
}),
|
||||
output_type: "markdown".to_string(),
|
||||
});
|
||||
|
||||
let broken_template = AnalysisTemplateSet {
|
||||
name: "E2E Broken Test".to_string(),
|
||||
modules,
|
||||
};
|
||||
|
||||
self.http_client.put(&format!("{}/api/v1/configs/templates/{}", GATEWAY_URL, broken_template_id))
|
||||
.json(&broken_template)
|
||||
.send().await?
|
||||
.error_for_status()?;
|
||||
|
||||
let symbol = "000001";
|
||||
let market = "CN";
|
||||
|
||||
let request_id = self.start_workflow(symbol, market, broken_template_id).await?;
|
||||
info!("Workflow started with broken template: {}", request_id);
|
||||
|
||||
match self.wait_for_completion(request_id, 30).await {
|
||||
Ok(true) => Err(anyhow!("Scenario E Failed: Workflow succeeded but should have failed")),
|
||||
Ok(false) => {
|
||||
info!("Scenario E Passed! Workflow failed as expected.");
|
||||
Ok(())
|
||||
},
|
||||
Err(e) => {
|
||||
warn!("Scenario E ended with error: {}", e);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn scenario_protocol_validation(&self) -> Result<()> {
|
||||
info!("=== Running Scenario: Protocol Validation (Reject/Timeout) ===");
|
||||
|
||||
// 1. Test Rejection
|
||||
let symbol = "TEST|reject";
|
||||
let market = "MOCK";
|
||||
let request_id = self.start_workflow(symbol, market, MOCK_TEMPLATE_ID).await?;
|
||||
info!("Started 'Reject' workflow: {}", request_id);
|
||||
|
||||
match self.wait_for_completion(request_id, 10).await {
|
||||
Ok(false) => info!("✅ Reject Scenario Passed (Workflow Failed as expected)"),
|
||||
Ok(true) => return Err(anyhow!("❌ Reject Scenario Failed (Workflow Succeeded unexpectedly)")),
|
||||
Err(e) => warn!("Reject Scenario ended with error check: {}", e),
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// 2. Test Timeout (No ACK)
|
||||
let symbol = "TEST|timeout_ack";
|
||||
let request_id = self.start_workflow(symbol, market, MOCK_TEMPLATE_ID).await?;
|
||||
info!("Started 'Timeout Ack' workflow: {}", request_id);
|
||||
|
||||
// Orchestrator timeout is 5s. We wait 10s.
|
||||
match self.wait_for_completion(request_id, 15).await {
|
||||
Ok(false) => info!("✅ Timeout Ack Scenario Passed (Workflow Failed as expected)"),
|
||||
Ok(true) => return Err(anyhow!("❌ Timeout Ack Scenario Failed (Workflow Succeeded unexpectedly)")),
|
||||
Err(e) => warn!("Timeout Ack Scenario ended with error check: {}", e),
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt::init();
|
||||
|
||||
let runner = TestRunner::new().await?;
|
||||
|
||||
runner.setup_test_environment().await?;
|
||||
|
||||
match runner.verify_registry_api().await {
|
||||
Ok(_) => info!("✅ Registry Verification PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Registry Verification FAILED: {}", e);
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// Run Protocol Validation
|
||||
match runner.scenario_protocol_validation().await {
|
||||
Ok(_) => info!("✅ Scenario Protocol Validation PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Scenario Protocol Validation FAILED: {}", e);
|
||||
// Don't return error yet, let other tests run? No, fail fast.
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// Run Mock Scenario First
|
||||
match runner.scenario_mock_workflow().await {
|
||||
Ok(_) => info!("✅ Scenario Mock PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Scenario Mock FAILED: {}", e);
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
match runner.scenario_a_happy_path().await {
|
||||
Ok(_) => info!("✅ Scenario A PASSED"),
|
||||
Err(e) => {
|
||||
// Soft fail if rate limited
|
||||
error!("❌ Scenario A FAILED: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
match runner.scenario_c_partial_failure().await {
|
||||
Ok(_) => info!("✅ Scenario C PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Scenario C FAILED: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
match runner.scenario_d_invalid_symbol().await {
|
||||
Ok(_) => info!("✅ Scenario D PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Scenario D FAILED: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
sleep(Duration::from_secs(2)).await;
|
||||
|
||||
match runner.scenario_e_analysis_failure().await {
|
||||
Ok(_) => info!("✅ Scenario E PASSED"),
|
||||
Err(e) => {
|
||||
error!("❌ Scenario E FAILED: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("🎉 All Scenarios Completed!");
|
||||
Ok(())
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user