owlen/crates/owlen-tui/tests/generation_tests.rs

use std::sync::{Arc, Mutex};
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use futures_util::stream;
use owlen_core::provider::{
    GenerateChunk, GenerateRequest, GenerateStream, ModelInfo, ModelProvider, ProviderMetadata,
    ProviderStatus, ProviderType,
};
use owlen_core::state::AppState;
use owlen_tui::app::{self, App, MessageState, messages::AppMessage};
use tokio::sync::mpsc;
use tokio::task::{JoinHandle, yield_now};
use tokio::time::advance;
use uuid::Uuid;

#[derive(Clone)]
struct StatusProvider {
    metadata: ProviderMetadata,
    status: Arc<Mutex<ProviderStatus>>,
    chunks: Arc<Vec<GenerateChunk>>,
}

impl StatusProvider {
    fn new(status: ProviderStatus, chunks: Vec<GenerateChunk>) -> Self {
        Self {
            metadata: ProviderMetadata::new("stub", "Stub", ProviderType::Local, false),
            status: Arc::new(Mutex::new(status)),
            chunks: Arc::new(chunks),
        }
    }

    fn set_status(&self, status: ProviderStatus) {
        *self.status.lock().unwrap() = status;
    }
}

#[async_trait]
impl ModelProvider for StatusProvider {
    fn metadata(&self) -> &ProviderMetadata {
        &self.metadata
    }

    async fn health_check(&self) -> Result<ProviderStatus, owlen_core::Error> {
        Ok(*self.status.lock().unwrap())
    }

    async fn list_models(&self) -> Result<Vec<ModelInfo>, owlen_core::Error> {
        Ok(vec![])
    }

    async fn generate_stream(
        &self,
        _request: GenerateRequest,
    ) -> Result<GenerateStream, owlen_core::Error> {
        let items = Arc::clone(&self.chunks);
        let stream_items = items.as_ref().clone();
        Ok(Box::pin(stream::iter(stream_items.into_iter().map(Ok))))
    }
}

#[derive(Default)]
struct RecordingState {
    started: bool,
    appended: bool,
    completed: bool,
    failed: bool,
    refreshed: bool,
    updated: bool,
    provider_status: Option<ProviderStatus>,
}

impl MessageState for RecordingState {
    fn start_generation(
        &mut self,
        _request_id: Uuid,
        _provider_id: &str,
        _request: &GenerateRequest,
    ) -> AppState {
        self.started = true;
        AppState::Running
    }

    fn append_chunk(&mut self, _request_id: Uuid, _chunk: &GenerateChunk) -> AppState {
        self.appended = true;
        AppState::Running
    }

    fn generation_complete(&mut self, _request_id: Uuid) -> AppState {
        self.completed = true;
        AppState::Running
    }

    fn generation_failed(&mut self, _request_id: Option<Uuid>, _message: &str) -> AppState {
        self.failed = true;
        AppState::Running
    }

    fn refresh_model_list(&mut self) -> AppState {
        self.refreshed = true;
        AppState::Running
    }

    fn update_model_list(&mut self) -> AppState {
        self.updated = true;
        AppState::Running
    }

    fn update_provider_status(&mut self, _provider_id: &str, status: ProviderStatus) -> AppState {
        self.provider_status = Some(status);
        AppState::Running
    }
}

#[tokio::test]
async fn start_and_abort_generation_manage_active_state() {
    let manager = Arc::new(owlen_core::provider::ProviderManager::default());
    let provider = StatusProvider::new(
        ProviderStatus::Available,
        vec![
            GenerateChunk::from_text("hello"),
            GenerateChunk::final_chunk(),
        ],
    );
    manager.register_provider(Arc::new(provider.clone())).await;
    let mut app = App::new(Arc::clone(&manager));

    let request_id = app
        .start_generation("stub", GenerateRequest::new("stub-model"))
        .expect("start generation");
    assert!(app.has_active_generation());
    assert_ne!(request_id, Uuid::nil());

    app.abort_active_generation();
    assert!(!app.has_active_generation());
}

#[test]
fn handle_message_dispatches_variants() {
    let manager = Arc::new(owlen_core::provider::ProviderManager::default());
    let mut app = App::new(Arc::clone(&manager));
    let mut state = RecordingState::default();
    let request_id = Uuid::new_v4();

    let _ = app.handle_message(
        &mut state,
        AppMessage::GenerateStart {
            request_id,
            provider_id: "stub".into(),
            request: GenerateRequest::new("stub"),
        },
    );
    let _ = app.handle_message(
        &mut state,
        AppMessage::GenerateChunk {
            request_id,
            chunk: GenerateChunk::from_text("chunk"),
        },
    );
    let _ = app.handle_message(&mut state, AppMessage::GenerateComplete { request_id });
    let _ = app.handle_message(
        &mut state,
        AppMessage::GenerateError {
            request_id: Some(request_id),
            message: "error".into(),
        },
    );
    let _ = app.handle_message(&mut state, AppMessage::ModelsRefresh);
    let _ = app.handle_message(&mut state, AppMessage::ModelsUpdated);
    let _ = app.handle_message(
        &mut state,
        AppMessage::ProviderStatus {
            provider_id: "stub".into(),
            status: ProviderStatus::Available,
        },
    );

    assert!(state.started);
    assert!(state.appended);
    assert!(state.completed);
    assert!(state.failed);
    assert!(state.refreshed);
    assert!(state.updated);
    assert!(matches!(
        state.provider_status,
        Some(ProviderStatus::Available)
    ));
}

#[tokio::test(start_paused = true)]
async fn background_worker_emits_status_changes() {
    let manager = Arc::new(owlen_core::provider::ProviderManager::default());
    let provider = StatusProvider::new(
        ProviderStatus::Unavailable,
        vec![GenerateChunk::final_chunk()],
    );
    manager.register_provider(Arc::new(provider.clone())).await;

    let (tx, mut rx) = mpsc::unbounded_channel();
    let worker: JoinHandle<()> = tokio::spawn(app::background_worker(Arc::clone(&manager), tx));

    provider.set_status(ProviderStatus::Available);
    advance(Duration::from_secs(31)).await;
    yield_now().await;

    if let Some(AppMessage::ProviderStatus { status, .. }) = rx.recv().await {
        assert!(matches!(status, ProviderStatus::Available));
    } else {
        panic!("expected provider status update");
    }

    worker.abort();
    let _ = worker.await;
    yield_now().await;
}