llm_worker_rs/worker/tests/common/mod.rs

//! テスト用共通ユーティリティ
//!
//! MockLlmClient、イベントレコーダー・プレイヤーを提供する

use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::Path;
use std::pin::Pin;
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use futures::Stream;
use serde::{Deserialize, Serialize};
use worker::llm_client::{ClientError, LlmClient, Request};
use worker_types::Event;

// =============================================================================
// Recorded Event Types
// =============================================================================

/// 記録されたSSEイベント
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecordedEvent {
    /// イベント受信からの経過時間 (ミリ秒)
    pub elapsed_ms: u64,
    /// SSEイベントタイプ
    pub event_type: String,
    /// SSEイベントデータ
    pub data: String,
}

/// セッションメタデータ
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMetadata {
    /// 記録開始タイムスタンプ (Unix epoch秒)
    pub timestamp: u64,
    /// モデル名
    pub model: String,
    /// リクエストの説明
    pub description: String,
}

// =============================================================================
// Event Recorder
// =============================================================================

/// SSEイベントレコーダー
///
/// 実際のAPIレスポンスを記録し、後でテストに使用できるようにする
#[allow(dead_code)]
pub struct EventRecorder {
    start_time: Instant,
    events: Vec<RecordedEvent>,
    metadata: SessionMetadata,
}

#[allow(dead_code)]
impl EventRecorder {
    /// 新しいレコーダーを作成
    pub fn new(model: impl Into<String>, description: impl Into<String>) -> Self {
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        Self {
            start_time: Instant::now(),
            events: Vec::new(),
            metadata: SessionMetadata {
                timestamp,
                model: model.into(),
                description: description.into(),
            },
        }
    }

    /// イベントを記録
    pub fn record(&mut self, event_type: &str, data: &str) {
        let elapsed = self.start_time.elapsed();
        self.events.push(RecordedEvent {
            elapsed_ms: elapsed.as_millis() as u64,
            event_type: event_type.to_string(),
            data: data.to_string(),
        });
    }

    /// 記録をファイルに保存
    ///
    /// フォーマット: JSONL (1行目: metadata, 2行目以降: events)
    pub fn save(&self, path: impl AsRef<Path>) -> std::io::Result<()> {
        let file = File::create(path)?;
        let mut writer = BufWriter::new(file);

        // メタデータを書き込み
        let metadata_json = serde_json::to_string(&self.metadata)?;
        writeln!(writer, "{}", metadata_json)?;

        // イベントを書き込み
        for event in &self.events {
            let event_json = serde_json::to_string(event)?;
            writeln!(writer, "{}", event_json)?;
        }

        writer.flush()?;
        Ok(())
    }

    /// 記録されたイベント数を取得
    pub fn event_count(&self) -> usize {
        self.events.len()
    }
}

// =============================================================================
// Event Player
// =============================================================================

/// SSEイベントプレイヤー
///
/// 記録されたイベントを読み込み、テストで使用する
#[allow(dead_code)]
pub struct EventPlayer {
    metadata: SessionMetadata,
    events: Vec<RecordedEvent>,
    current_index: usize,
}

#[allow(dead_code)]
impl EventPlayer {
    /// ファイルから読み込み
    pub fn load(path: impl AsRef<Path>) -> std::io::Result<Self> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let mut lines = reader.lines();

        // メタデータを読み込み
        let metadata_line = lines
            .next()
            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "Empty file"))??;
        let metadata: SessionMetadata = serde_json::from_str(&metadata_line)?;

        // イベントを読み込み
        let mut events = Vec::new();
        for line in lines {
            let line = line?;
            if !line.is_empty() {
                let event: RecordedEvent = serde_json::from_str(&line)?;
                events.push(event);
            }
        }

        Ok(Self {
            metadata,
            events,
            current_index: 0,
        })
    }

    /// メタデータを取得
    pub fn metadata(&self) -> &SessionMetadata {
        &self.metadata
    }

    /// 全イベントを取得
    pub fn events(&self) -> &[RecordedEvent] {
        &self.events
    }

    /// イベント数を取得
    pub fn event_count(&self) -> usize {
        self.events.len()
    }

    /// 次のイベントを取得（Iterator的に使用）
    pub fn next_event(&mut self) -> Option<&RecordedEvent> {
        if self.current_index < self.events.len() {
            let event = &self.events[self.current_index];
            self.current_index += 1;
            Some(event)
        } else {
            None
        }
    }

    /// インデックスをリセット
    pub fn reset(&mut self) {
        self.current_index = 0;
    }

    /// 全イベントをworker_types::Eventとしてパースして取得
    pub fn parse_events(&self) -> Vec<Event> {
        self.events
            .iter()
            .filter_map(|recorded| serde_json::from_str(&recorded.data).ok())
            .collect()
    }
}

// =============================================================================
// MockLlmClient
// =============================================================================

/// テスト用のモックLLMクライアント
///
/// 事前に定義されたイベントシーケンスをストリームとして返す。
/// fixtureファイルからロードすることも、直接イベントを渡すこともできる。
pub struct MockLlmClient {
    events: Vec<Event>,
}

impl MockLlmClient {
    /// イベントリストから直接作成
    pub fn new(events: Vec<Event>) -> Self {
        Self { events }
    }

    /// fixtureファイルからロード
    pub fn from_fixture(path: impl AsRef<Path>) -> std::io::Result<Self> {
        let player = EventPlayer::load(path)?;
        let events = player.parse_events();
        Ok(Self { events })
    }

    /// 保持しているイベント数を取得
    pub fn event_count(&self) -> usize {
        self.events.len()
    }
}

#[async_trait]
impl LlmClient for MockLlmClient {
    async fn stream(
        &self,
        _request: Request,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
        let events = self.events.clone();
        let stream = futures::stream::iter(events.into_iter().map(Ok));
        Ok(Box::pin(stream))
    }
}