102 lines
3.2 KiB
Rust
102 lines
3.2 KiB
Rust
use std::time::{Duration, Instant};
|
|
|
|
use crank_core::{ExecutionFailure, ExecutionOrigin, InvocationSource, WorkspaceId};
|
|
use crank_runtime::{
|
|
ExecutionAuthorization, ResolvedAuth, RuntimeExecutionRequest, RuntimeExecutor,
|
|
RuntimeOperation, RuntimeRequestContext,
|
|
};
|
|
use serde_json::Value;
|
|
|
|
#[allow(dead_code)]
|
|
pub trait RuntimeExecutorTestExt {
|
|
async fn execute(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
) -> Result<Value, ExecutionFailure>;
|
|
|
|
async fn execute_with_context(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
context: Option<&RuntimeRequestContext>,
|
|
) -> Result<Value, ExecutionFailure>;
|
|
|
|
async fn execute_with_auth_and_context(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
resolved_auth: Option<&ResolvedAuth>,
|
|
context: Option<&RuntimeRequestContext>,
|
|
) -> Result<Value, ExecutionFailure>;
|
|
}
|
|
|
|
impl RuntimeExecutorTestExt for RuntimeExecutor {
|
|
async fn execute(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
) -> Result<Value, ExecutionFailure> {
|
|
self.execute_with_auth_and_context(operation, input, None, None)
|
|
.await
|
|
}
|
|
|
|
async fn execute_with_context(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
context: Option<&RuntimeRequestContext>,
|
|
) -> Result<Value, ExecutionFailure> {
|
|
self.execute_with_auth_and_context(operation, input, None, context)
|
|
.await
|
|
}
|
|
|
|
async fn execute_with_auth_and_context(
|
|
&self,
|
|
operation: &RuntimeOperation,
|
|
input: &Value,
|
|
resolved_auth: Option<&ResolvedAuth>,
|
|
context: Option<&RuntimeRequestContext>,
|
|
) -> Result<Value, ExecutionFailure> {
|
|
let mut context = context.cloned().unwrap_or_else(|| {
|
|
RuntimeRequestContext::from_correlation(&crank_core::CorrelationContext::generate())
|
|
});
|
|
if context.metering_context().is_none() {
|
|
context = context.with_metering_context(
|
|
WorkspaceId::new("ws_test"),
|
|
None,
|
|
InvocationSource::AdminTestRun,
|
|
);
|
|
}
|
|
let metering = context.metering_context().expect("test metering context");
|
|
let origin = if metering.agent_id.is_some() {
|
|
ExecutionOrigin::AgentSnapshot
|
|
} else {
|
|
ExecutionOrigin::AdminDraft
|
|
};
|
|
let request = RuntimeExecutionRequest::try_new(
|
|
&metering.workspace_id,
|
|
origin,
|
|
metering.agent_id.as_ref(),
|
|
operation,
|
|
input,
|
|
ExecutionAuthorization::Authorized,
|
|
resolved_auth,
|
|
&context,
|
|
Instant::now() + Duration::from_secs(30),
|
|
)
|
|
.map_err(|_| {
|
|
crank_core::ExecutionFailure::new(
|
|
crank_core::ExecutionErrorCode::RuntimeInternal,
|
|
crank_core::CorrelationContext::new(
|
|
context.request_id.clone(),
|
|
context.trace_context.clone(),
|
|
),
|
|
)
|
|
})?;
|
|
self.execute_outcome(request)
|
|
.await
|
|
.map(|success| success.output)
|
|
}
|
|
}
|