oiupoyt/brokecrew
1
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

test: add comprehensive integration tests and WireMock test harness

Browse source
This commit is contained in:
katR 2026-09-26 13:55:12 +05:30
parent fe3fc74e1c
commit b7b88c6daa
14 changed files with 706 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
src/budget.rs
View file

@ -163,8 +163,7 @@ impl BudgetTracker {
last_day_utc: self.last_day_utc.clone(),
recent_calls: self.recent_calls.clone(),
};
let serialized = serde_json::to_string_pretty(&state)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
let serialized = serde_json::to_string_pretty(&state).map_err(std::io::Error::other)?;
// Atomic write via tempfile in the same parent directory
let parent_dir = path.parent().unwrap_or_else(|| Path::new("."));

13
src/cli.rs
View file

@ -15,10 +15,19 @@ pub struct Cli {
#[arg(long, global = true, help = "Path to personas.toml file")]
pub personas_file: Option<PathBuf>,
#[arg(long, global = true, env = "OPENROUTER_API_KEY", help = "OpenRouter API Key (overrides keyring & env)")]
#[arg(
long,
global = true,
env = "OPENROUTER_API_KEY",
help = "OpenRouter API Key (overrides keyring & env)"
)]
pub api_key: Option<String>,
#[arg(long, global = true, help = "Output structured logs in JSON format to stderr")]
#[arg(
long,
global = true,
help = "Output structured logs in JSON format to stderr"
)]
pub json_logs: bool,
#[arg(long, short, global = true, help = "Enable verbose debug logs")]

45
src/client.rs
View file

@ -136,7 +136,11 @@ impl OpenRouterClient {
None
}
fn check_status(status: reqwest::StatusCode, body: &str, headers: &reqwest::header::HeaderMap) -> Result<(), ClientError> {
fn check_status(
status: reqwest::StatusCode,
body: &str,
headers: &reqwest::header::HeaderMap,
) -> Result<(), ClientError> {
match status.as_u16() {
200..=299 => Ok(()),
401 => Err(ClientError::Unauthorized),
@ -149,7 +153,9 @@ impl OpenRouterClient {
.unwrap_or(5);
Err(ClientError::RateLimited(retry_after))
}
404 => Err(ClientError::ModelUnavailable(format!("Model not found (404): {body}"))),
404 => Err(ClientError::ModelUnavailable(format!(
"Model not found (404): {body}"
))),
s if s >= 500 => Err(ClientError::ServerError {
status: s,
message: body.to_string(),
@ -223,7 +229,8 @@ impl OpenRouterClient {
#[async_trait]
impl ModelClient for OpenRouterClient {
async fn complete(&self, model: &str, system: &str, user: &str) -> Result<String, ClientError> {
self.complete_with_params(model, system, user, None, None).await
self.complete_with_params(model, system, user, None, None)
.await
}
async fn complete_stream(
@ -276,7 +283,7 @@ impl ModelClient for OpenRouterClient {
let mut buffer = String::new();
while let Some(item) = stream.next().await {
let bytes = item.map_err(|e| ClientError::Network(e))?;
let bytes = item.map_err(ClientError::Network)?;
let chunk_str = String::from_utf8_lossy(&bytes);
buffer.push_str(&chunk_str);
@ -331,27 +338,47 @@ mod tests {
#[test]
fn test_check_status_mappings() {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert("retry-after", reqwest::header::HeaderValue::from_static("12"));
headers.insert(
"retry-after",
reqwest::header::HeaderValue::from_static("12"),
);
let err_401 = OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers).unwrap_err();
let err_401 =
OpenRouterClient::check_status(reqwest::StatusCode::UNAUTHORIZED, "bad key", &headers)
.unwrap_err();
match err_401 {
ClientError::Unauthorized => (),
_ => panic!("Expected Unauthorized, got {:?}", err_401),
}
let err_402 = OpenRouterClient::check_status(reqwest::StatusCode::PAYMENT_REQUIRED, "low balance", &headers).unwrap_err();
let err_402 = OpenRouterClient::check_status(
reqwest::StatusCode::PAYMENT_REQUIRED,
"low balance",
&headers,
)
.unwrap_err();
match err_402 {
ClientError::PaymentRequired => (),
_ => panic!("Expected PaymentRequired, got {:?}", err_402),
}
let err_429 = OpenRouterClient::check_status(reqwest::StatusCode::TOO_MANY_REQUESTS, "rate limited", &headers).unwrap_err();
let err_429 = OpenRouterClient::check_status(
reqwest::StatusCode::TOO_MANY_REQUESTS,
"rate limited",
&headers,
)
.unwrap_err();
match err_429 {
ClientError::RateLimited(secs) => assert_eq!(secs, 12),
_ => panic!("Expected RateLimited(12), got {:?}", err_429),
}
let err_503 = OpenRouterClient::check_status(reqwest::StatusCode::SERVICE_UNAVAILABLE, "outage", &headers).unwrap_err();
let err_503 = OpenRouterClient::check_status(
reqwest::StatusCode::SERVICE_UNAVAILABLE,
"outage",
&headers,
)
.unwrap_err();
match err_503 {
ClientError::ServerError { status, .. } => assert_eq!(status, 503),
_ => panic!("Expected ServerError 503, got {:?}", err_503),

2
src/config.rs
View file

@ -235,6 +235,6 @@ mod tests {
fn test_bundled_personas_valid() {
let parsed: PersonasConfigFile = toml::from_str(DEFAULT_PERSONAS_TOML).unwrap();
assert!(parsed.personas.len() >= 3);
assert!(parsed.pipelines.len() >= 1);
assert!(!parsed.pipelines.is_empty());
}
}

82
src/main.rs
View file

@ -103,7 +103,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
AppConfig::save_api_key_to_keyring(&api_key)
.map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?;
println!("{}", "✔ OpenRouter API key saved securely in OS credential store.".bright_green().bold());
println!(
"{}",
"✔ OpenRouter API key saved securely in OS credential store."
.bright_green()
.bold()
);
return Ok(());
}
},
@ -115,7 +120,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let badge = get_persona_color_tag(&persona.name);
println!("\n {} - {}", badge, persona.system_prompt.dimmed());
if !persona.trigger_words.is_empty() {
println!(" Triggers: {}", persona.trigger_words.join(", ").italic());
println!(
" Triggers: {}",
persona.trigger_words.join(", ").italic()
);
}
println!(" Candidates:");
for (i, c) in persona.candidates.iter().enumerate() {
@ -128,7 +136,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
} else {
String::new()
};
println!(" {}. {} ({}){}", i + 1, c.id.bold(), tier_tag, stealth_tag);
println!(
" {}. {} ({}){}",
i + 1,
c.id.bold(),
tier_tag,
stealth_tag
);
}
}
@ -146,7 +160,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}
})
.collect();
println!(" • {} ➔ {}", pipe.name.bright_cyan().bold(), steps_str.join(" ➔ "));
println!(
" • {} ➔ {}",
pipe.name.bright_cyan().bold(),
steps_str.join(" ➔ ")
);
}
}
println!();
@ -171,20 +189,39 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
};
if model_count > 0 {
println!(" Live catalog: {} available models cached", model_count.to_string().bright_green().bold());
println!(
" Live catalog: {} available models cached",
model_count.to_string().bright_green().bold()
);
}
match budget_tracker.check_budget() {
Ok(status) => {
let pct = (status.daily_requests * 100) / status.daily_limit;
let daily_color = if pct >= 80 {
format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).yellow().bold().to_string()
format!(
"{}/{} ({}%)",
status.daily_requests, status.daily_limit, pct
)
.yellow()
.bold()
.to_string()
} else {
format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).green().bold().to_string()
format!(
"{}/{} ({}%)",
status.daily_requests, status.daily_limit, pct
)
.green()
.bold()
.to_string()
};
println!(" Account tier: {:?}", config.credit_tier);
println!(" Rolling 1-min requests: {}/{}", status.minute_requests.to_string().bold(), status.minute_limit);
println!(
" Rolling 1-min requests: {}/{}",
status.minute_requests.to_string().bold(),
status.minute_limit
);
println!(" Daily requests today: {}", daily_color);
if status.approaching_daily_limit {
println!(" {}", "⚠ Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold());
@ -204,7 +241,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
None => {
// Default to chat if no subcommand passed
run_chat(ChatArgs::default(), &config, &mut budget_tracker, model_registry).await?;
run_chat(
ChatArgs::default(),
&config,
&mut budget_tracker,
model_registry,
)
.await?;
}
}
@ -261,8 +304,14 @@ async fn run_chat(
}
// Otherwise, run interactive REPL loop
println!("\n{}", "BrokeCrew Interactive Team Chat".bright_cyan().bold());
println!("{}", "Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed());
println!(
"\n{}",
"BrokeCrew Interactive Team Chat".bright_cyan().bold()
);
println!(
"{}",
"Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed()
);
let stdin = io::stdin();
let mut reader = stdin.lock();
@ -375,14 +424,20 @@ async fn execute_prompt(
(dec.persona, Some(info))
}
Err(e) => {
eprintln!("Router warning: {}, falling back to {}", e, config.default_persona);
eprintln!(
"Router warning: {}, falling back to {}",
e, config.default_persona
);
(config.default_persona.clone(), None)
}
}
};
if let Some(info) = route_info {
eprintln!("{}", format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed());
eprintln!(
"{}",
format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed()
);
}
let (chunk_tx, mut chunk_rx) = mpsc::channel::<String>(100);
@ -435,6 +490,7 @@ async fn execute_prompt(
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn run_pipeline(
pipeline: &PipelineConfig,
config: &AppConfig,

12
src/persona.rs
View file

@ -49,7 +49,15 @@ impl PersonaRegistry {
pub fn new(personas: Vec<Persona>) -> Self {
Self { personas }
}
}
impl From<Vec<Persona>> for PersonaRegistry {
fn from(personas: Vec<Persona>) -> Self {
Self::new(personas)
}
}
impl PersonaRegistry {
pub fn get(&self, name: &str) -> Option<&Persona> {
self.personas.iter().find(|p| p.name == name)
}
@ -152,6 +160,8 @@ mod tests {
assert_eq!(paid_only[0].id, "gemini-paid");
// Non-existent persona returns error
assert!(registry.get_eligible_candidates("unknown", false, false).is_err());
assert!(registry
.get_eligible_candidates("unknown", false, false)
.is_err());
}
}

27
src/pipeline.rs
View file

@ -45,6 +45,7 @@ pub struct PipelineExecutor;
impl PipelineExecutor {
/// Executes a single persona query with candidate fallback walk
#[allow(clippy::too_many_arguments)]
pub async fn dispatch_persona(
persona_name: &str,
persona_registry: &PersonaRegistry,
@ -55,9 +56,9 @@ impl PipelineExecutor {
require_paid: bool,
chunk_sender: Option<mpsc::Sender<String>>,
) -> Result<DispatchResult, ClientError> {
let persona = persona_registry
.get(persona_name)
.ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")))?;
let persona = persona_registry.get(persona_name).ok_or_else(|| {
ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found"))
})?;
// Filter candidates against live registry and options
let eligible = model_registry
@ -68,7 +69,11 @@ impl PipelineExecutor {
eligible
} else {
// If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates
persona.candidates.iter().filter(|c| allow_stealth || !c.is_stealth).collect()
persona
.candidates
.iter()
.filter(|c| allow_stealth || !c.is_stealth)
.collect()
};
if candidates_to_try.is_empty() {
@ -166,6 +171,7 @@ impl PipelineExecutor {
}
/// Executes an entire pipeline sequentially, passing outputs between steps
#[allow(clippy::too_many_arguments)]
pub async fn execute(
pipeline: &PipelineConfig,
persona_registry: &PersonaRegistry,
@ -181,9 +187,9 @@ impl PipelineExecutor {
let total_steps = pipeline.steps.len();
for (step_idx, step) in pipeline.steps.iter().enumerate() {
let persona = persona_registry
.get(&step.persona)
.ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)))?;
let persona = persona_registry.get(&step.persona).ok_or_else(|| {
ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona))
})?;
// Prepare prompt for this step
let step_prompt = if step.pass_previous_output && !previous_output.is_empty() {
@ -286,7 +292,12 @@ mod tests {
#[async_trait]
impl ModelClient for MockPipelineClient {
async fn complete(&self, model: &str, _system: &str, user: &str) -> Result<String, ClientError> {
async fn complete(
&self,
model: &str,
_system: &str,
user: &str,
) -> Result<String, ClientError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(format!("Output from {} for: {}", model, user))
}

16
src/registry.rs
View file

@ -98,7 +98,10 @@ impl ModelRegistry {
.map_err(|e| format!("Network error fetching models catalog: {e}"))?;
if !res.status().is_success() {
return Err(format!("OpenRouter models endpoint returned HTTP {}", res.status()));
return Err(format!(
"OpenRouter models endpoint returned HTTP {}",
res.status()
));
}
let body = res
@ -184,10 +187,17 @@ mod tests {
]
}"#;
let count = registry.update_from_json(json_data).await.expect("update failed");
let count = registry
.update_from_json(json_data)
.await
.expect("update failed");
assert_eq!(count, 2);
assert!(registry.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free").await);
assert!(
registry
.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free")
.await
);
assert!(!registry.is_model_available("unlisted/model").await);
let candidates = vec![

20
src/router/keyword.rs
View file

@ -70,13 +70,24 @@ mod tests {
Persona {
name: "coder".into(),
system_prompt: "coder prompt".into(),
trigger_words: vec!["build".into(), "code".into(), "script".into(), "website".into(), "api".into()],
trigger_words: vec![
"build".into(),
"code".into(),
"script".into(),
"website".into(),
"api".into(),
],
candidates: vec![],
},
Persona {
name: "reviewer".into(),
system_prompt: "reviewer prompt".into(),
trigger_words: vec!["review".into(), "audit".into(), "critique".into(), "lint".into()],
trigger_words: vec![
"review".into(),
"audit".into(),
"critique".into(),
"lint".into(),
],
candidates: vec![],
},
Persona {
@ -97,7 +108,10 @@ mod tests {
("write a python script to parse logs", Some("coder")),
("please review this pull request for bugs", Some("reviewer")),
("audit our authentication system", Some("reviewer")),
("design a high scale distributed cache architecture", Some("architect")),
(
"design a high scale distributed cache architecture",
Some("architect"),
),
("hello, what is the weather today?", None),
("how are you doing?", None),
];

51
src/router/llm.rs
View file

@ -15,9 +15,15 @@ impl LlmRouter {
known_personas.join(", ")
);
let user_prompt = format!("Classify this request to the most appropriate persona:\n\"{}\"", prompt);
let user_prompt = format!(
"Classify this request to the most appropriate persona:\n\"{}\"",
prompt
);
let raw_output = match client.complete(router_model, &system_prompt, &user_prompt).await {
let raw_output = match client
.complete(router_model, &system_prompt, &user_prompt)
.await
{
Ok(out) => out,
Err(e) => {
tracing::warn!("Router LLM call failed, falling back to default persona: {e}");
@ -26,9 +32,15 @@ impl LlmRouter {
};
let cleaned = Self::sanitize_output(&raw_output);
if known_personas.iter().any(|p| p.eq_ignore_ascii_case(&cleaned)) {
if known_personas
.iter()
.any(|p| p.eq_ignore_ascii_case(&cleaned))
{
// Find matched casing
if let Some(matched) = known_personas.iter().find(|p| p.eq_ignore_ascii_case(&cleaned)) {
if let Some(matched) = known_personas
.iter()
.find(|p| p.eq_ignore_ascii_case(&cleaned))
{
return Ok(matched.clone());
}
}
@ -94,7 +106,12 @@ mod tests {
#[async_trait]
impl ModelClient for MockRouterClient {
async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result<String, ClientError> {
async fn complete(
&self,
_model: &str,
_system: &str,
_user: &str,
) -> Result<String, ClientError> {
Ok(self.reply.clone())
}
@ -117,7 +134,13 @@ mod tests {
let client_valid = MockRouterClient {
reply: "coder".into(),
};
let res = LlmRouter::classify(&client_valid, "router-model", &known, "general", "build something")
let res = LlmRouter::classify(
&client_valid,
"router-model",
&known,
"general",
"build something",
)
.await
.unwrap();
assert_eq!(res, "coder");
@ -126,7 +149,13 @@ mod tests {
let client_messy = MockRouterClient {
reply: " \"REVIEWER\" \n".into(),
};
let res = LlmRouter::classify(&client_messy, "router-model", &known, "general", "audit this")
let res = LlmRouter::classify(
&client_messy,
"router-model",
&known,
"general",
"audit this",
)
.await
.unwrap();
assert_eq!(res, "reviewer");
@ -135,7 +164,13 @@ mod tests {
let client_hallucinated = MockRouterClient {
reply: "I believe a poet would be best for this query".into(),
};
let res = LlmRouter::classify(&client_hallucinated, "router-model", &known, "general", "write poem")
let res = LlmRouter::classify(
&client_hallucinated,
"router-model",
&known,
"general",
"write poem",
)
.await
.unwrap();
assert_eq!(res, "general");

23
src/router/mod.rs
View file

@ -87,7 +87,12 @@ mod tests {
#[async_trait]
impl ModelClient for DummyClient {
async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result<String, ClientError> {
async fn complete(
&self,
_model: &str,
_system: &str,
_user: &str,
) -> Result<String, ClientError> {
Ok("general".into())
}
async fn complete_stream(
@ -118,17 +123,27 @@ mod tests {
},
];
let registry = PersonaRegistry::new(personas);
let router = Router::new("general".into(), "meta-llama/llama-3.2-3b-instruct:free".into(), 7);
let router = Router::new(
"general".into(),
"meta-llama/llama-3.2-3b-instruct:free".into(),
7,
);
let client = DummyClient;
// Matches Stage A (build) without escalation
let dec1 = router.route(&registry, &client, "build a script", false).await.unwrap();
let dec1 = router
.route(&registry, &client, "build a script", false)
.await
.unwrap();
assert_eq!(dec1.persona, "coder");
assert_eq!(dec1.stage_used, RouteStage::StageAKeyword);
assert!(!dec1.escalate_to_paid);
// Matches Stage A with forced escalation
let dec2 = router.route(&registry, &client, "build a script", true).await.unwrap();
let dec2 = router
.route(&registry, &client, "build a script", true)
.await
.unwrap();
assert_eq!(dec2.persona, "coder");
assert!(dec2.escalate_to_paid);
}

174
tests/fallback_tests.rs Normal file
View file

@ -0,0 +1,174 @@
use async_trait::async_trait;
use brokecrew::client::{ClientError, ModelClient};
use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier};
use brokecrew::pipeline::PipelineExecutor;
use brokecrew::registry::ModelRegistry;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;
struct FaultInjectingClient {
candidate1_error: Option<ClientError>,
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl ModelClient for FaultInjectingClient {
async fn complete(&self, model: &str, system: &str, user: &str) -> Result<String, ClientError> {
self.calls.fetch_add(1, Ordering::SeqCst);
if model == "candidate-1" {
if let Some(ref err) = self.candidate1_error {
return match err {
ClientError::RateLimited(s) => Err(ClientError::RateLimited(*s)),
ClientError::ServerError { status, message } => Err(ClientError::ServerError {
status: *status,
message: message.clone(),
}),
ClientError::Unauthorized => Err(ClientError::Unauthorized),
ClientError::PaymentRequired => Err(ClientError::PaymentRequired),
_ => Err(ClientError::ModelUnavailable("simulated".into())),
};
}
}
Ok(format!(
"Success from {model} for prompt '{user}' (system: {system})"
))
}
async fn complete_stream(
&self,
model: &str,
system: &str,
user: &str,
chunk_sender: mpsc::Sender<String>,
) -> Result<String, ClientError> {
let reply = self.complete(model, system, user).await?;
let _ = chunk_sender.send(reply.clone()).await;
Ok(reply)
}
}
fn create_resilience_persona() -> Persona {
Persona {
name: "resilient-coder".into(),
system_prompt: "You are coder".into(),
trigger_words: vec!["code".into()],
candidates: vec![
ModelCandidate {
id: "candidate-1".into(),
tier: Tier::Free,
is_stealth: false,
},
ModelCandidate {
id: "candidate-2".into(),
tier: Tier::Free,
is_stealth: false,
},
],
}
}
#[tokio::test]
async fn test_candidate_cascade_on_503_service_unavailable() {
let persona = create_resilience_persona();
let persona_reg = PersonaRegistry::new(vec![persona]);
let model_reg = ModelRegistry::new();
let calls = Arc::new(AtomicUsize::new(0));
let client = FaultInjectingClient {
candidate1_error: Some(ClientError::ServerError {
status: 503,
message: "Overloaded".into(),
}),
calls: calls.clone(),
};
let result = PipelineExecutor::dispatch_persona(
"resilient-coder",
&persona_reg,
&model_reg,
&client,
"write rust code",
false,
false,
None,
)
.await
.expect("fallback to candidate-2 should succeed");
assert_eq!(result.persona, "resilient-coder");
assert_eq!(result.model_used, "candidate-2");
assert!(result.fell_back, "Result must indicate fallback occurred");
assert_eq!(result.attempts.len(), 1);
assert_eq!(result.attempts[0].model_id, "candidate-1");
assert_eq!(
calls.load(Ordering::SeqCst),
2,
"Both candidates must have been called"
);
}
#[tokio::test]
async fn test_candidate_cascade_on_429_rate_limited() {
let persona = create_resilience_persona();
let persona_reg = PersonaRegistry::new(vec![persona]);
let model_reg = ModelRegistry::new();
let calls = Arc::new(AtomicUsize::new(0));
let client = FaultInjectingClient {
candidate1_error: Some(ClientError::RateLimited(15)),
calls: calls.clone(),
};
let result = PipelineExecutor::dispatch_persona(
"resilient-coder",
&persona_reg,
&model_reg,
&client,
"write python code",
false,
false,
None,
)
.await
.expect("fallback to candidate-2 should succeed");
assert_eq!(result.model_used, "candidate-2");
assert!(result.fell_back);
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_fail_fast_on_401_unauthorized() {
let persona = create_resilience_persona();
let persona_reg = PersonaRegistry::new(vec![persona]);
let model_reg = ModelRegistry::new();
let calls = Arc::new(AtomicUsize::new(0));
let client = FaultInjectingClient {
candidate1_error: Some(ClientError::Unauthorized),
calls: calls.clone(),
};
let err = PipelineExecutor::dispatch_persona(
"resilient-coder",
&persona_reg,
&model_reg,
&client,
"test query",
false,
false,
None,
)
.await
.unwrap_err();
match err {
ClientError::Unauthorized => (),
_ => panic!("Expected Unauthorized fail-fast, got {:?}", err),
}
// Must NOT call candidate-2 when 401 Unauthorized occurs!
assert_eq!(calls.load(Ordering::SeqCst), 1);
}

148
tests/router_tests.rs Normal file
View file

@ -0,0 +1,148 @@
use async_trait::async_trait;
use brokecrew::client::{ClientError, ModelClient};
use brokecrew::persona::{ModelCandidate, Persona, PersonaRegistry, Tier};
use brokecrew::router::{RouteDecision, RouteStage, Router};
use tokio::sync::mpsc;
struct TestRouterClient {
llm_choice: String,
}
#[async_trait]
impl ModelClient for TestRouterClient {
async fn complete(
&self,
_model: &str,
_system: &str,
_user: &str,
) -> Result<String, ClientError> {
Ok(self.llm_choice.clone())
}
async fn complete_stream(
&self,
model: &str,
system: &str,
user: &str,
_chunk_sender: mpsc::Sender<String>,
) -> Result<String, ClientError> {
self.complete(model, system, user).await
}
}
fn setup_test_registry() -> PersonaRegistry {
vec![
Persona {
name: "coder".into(),
system_prompt: "You write code".into(),
trigger_words: vec![
"build".into(),
"code".into(),
"script".into(),
"website".into(),
"api".into(),
],
candidates: vec![ModelCandidate {
id: "qwen-free".into(),
tier: Tier::Free,
is_stealth: false,
}],
},
Persona {
name: "reviewer".into(),
system_prompt: "You review code".into(),
trigger_words: vec!["review".into(), "audit".into(), "critique".into()],
candidates: vec![ModelCandidate {
id: "llama-free".into(),
tier: Tier::Free,
is_stealth: false,
}],
},
Persona {
name: "general".into(),
system_prompt: "You are general assistant".into(),
trigger_words: vec!["help".into(), "hello".into()],
candidates: vec![ModelCandidate {
id: "general-free".into(),
tier: Tier::Free,
is_stealth: false,
}],
},
]
.into()
}
#[tokio::test]
async fn test_table_driven_classification() {
let registry = setup_test_registry();
let client = TestRouterClient {
llm_choice: "reviewer".into(),
};
let router = Router::new("general".into(), "router-model".into(), 7);
let cases = vec![
(
"build me a proxy checker site",
"coder",
RouteStage::StageAKeyword,
false,
),
(
"please audit this auth module",
"reviewer",
RouteStage::StageAKeyword,
false,
),
(
"how does photosynthesis work?",
"reviewer",
RouteStage::StageBLlm,
false,
), // LLM choice returned "reviewer"
];
for (prompt, expected_persona, expected_stage, expected_escalate) in cases {
let decision: RouteDecision = router
.route(&registry, &client, prompt, false)
.await
.expect("routing should succeed");
assert_eq!(
decision.persona, expected_persona,
"Prompt failed: '{}'",
prompt
);
assert_eq!(
decision.stage_used, expected_stage,
"Stage mismatch for: '{}'",
prompt
);
assert_eq!(
decision.escalate_to_paid, expected_escalate,
"Escalate mismatch for: '{}'",
prompt
);
}
}
#[tokio::test]
async fn test_complexity_escalation_trigger() {
let registry = setup_test_registry();
let client = TestRouterClient {
llm_choice: "coder".into(),
};
let router = Router::new("general".into(), "router-model".into(), 7);
let heavy_prompt = "Refactor all microservices across the entire project with database migration and distributed architecture:\n```rust\nfn main() {}\n```";
let decision = router
.route(&registry, &client, heavy_prompt, false)
.await
.expect("routing should succeed");
assert_eq!(decision.persona, "coder");
assert!(
decision.escalate_to_paid,
"Heavy prompt must trigger escalation"
);
assert!(decision.complexity_score >= 7);
}

138
tests/wiremock_client.rs Normal file
View file

@ -0,0 +1,138 @@
use brokecrew::client::{ClientError, ModelClient, OpenRouterClient};
use brokecrew::registry::ModelRegistry;
use tokio::sync::mpsc;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[tokio::test]
async fn test_wiremock_sse_streaming() {
let mock_server = MockServer::start().await;
let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"from \"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\"WireMock!\"}}]}\n\ndata: [DONE]\n\n";
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(sse_body),
)
.mount(&mock_server)
.await;
let client = OpenRouterClient::with_base_url("test-key".into(), mock_server.uri());
let (tx, mut rx) = mpsc::channel::<String>(10);
let collect_task = tokio::spawn(async move {
let mut chunks = Vec::new();
while let Some(c) = rx.recv().await {
chunks.push(c);
}
chunks
});
let full_reply = client
.complete_stream("mock-model", "system", "user", tx)
.await
.expect("streaming should succeed");
let chunks = collect_task.await.unwrap();
assert_eq!(full_reply, "Hello from WireMock!");
assert_eq!(chunks, vec!["Hello ", "from ", "WireMock!"]);
}
#[tokio::test]
async fn test_wiremock_401_unauthorized() {
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(ResponseTemplate::new(401).set_body_string("Invalid API key"))
.mount(&mock_server)
.await;
let client = OpenRouterClient::with_base_url("bad-key".into(), mock_server.uri());
let err = client
.complete("mock-model", "sys", "user")
.await
.unwrap_err();
match err {
ClientError::Unauthorized => (),
_ => panic!("Expected ClientError::Unauthorized, got {:?}", err),
}
}
#[tokio::test]
async fn test_wiremock_402_payment_required() {
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(ResponseTemplate::new(402).set_body_string("Account balance depleted"))
.mount(&mock_server)
.await;
let client = OpenRouterClient::with_base_url("key".into(), mock_server.uri());
let err = client
.complete("mock-model", "sys", "user")
.await
.unwrap_err();
match err {
ClientError::PaymentRequired => (),
_ => panic!("Expected ClientError::PaymentRequired, got {:?}", err),
}
}
#[tokio::test]
async fn test_wiremock_models_catalog_fetch() {
let mock_server = MockServer::start().await;
let catalog_json = r#"{
"data": [
{
"id": "qwen/qwen-2.5-coder-32b-instruct:free",
"name": "Qwen 2.5 Coder",
"context_length": 32768,
"pricing": { "prompt": "0", "completion": "0" }
},
{
"id": "meta-llama/llama-3.3-70b-instruct:free",
"name": "Llama 3.3 70B",
"context_length": 131072,
"pricing": { "prompt": "0", "completion": "0" }
}
]
}"#;
Mock::given(method("GET"))
.and(path("/models"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "application/json")
.set_body_string(catalog_json),
)
.mount(&mock_server)
.await;
let registry = ModelRegistry::new();
let count = registry
.fetch_live(&mock_server.uri())
.await
.expect("catalog fetch failed");
assert_eq!(count, 2);
assert!(
registry
.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free")
.await
);
assert!(
registry
.is_model_available("meta-llama/llama-3.3-70b-instruct:free")
.await
);
assert!(!registry.is_model_available("non-existent-model").await);
}