brokecrew/tests/router_tests.rs

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);
}