Skip to main content

research_master/sources/
mock.rs

1//! Mock source for testing purposes.
2
3use async_trait::async_trait;
4use std::sync::Mutex;
5
6use crate::models::{Paper, SearchQuery, SearchResponse, SourceType};
7use crate::sources::{Source, SourceCapabilities, SourceError};
8
9/// A mock source for testing that returns predefined responses.
10#[derive(Debug, Default, Clone)]
11pub struct MockSource {
12    // Note: Clone won't deep clone the mutex contents, but for testing
13    // this is acceptable as tests typically use the original reference.
14    search_response: std::sync::Arc<Mutex<Option<SearchResponse>>>,
15    search_count: std::sync::Arc<Mutex<usize>>,
16}
17
18impl MockSource {
19    /// Create a new mock source.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Set the search response to return.
25    pub fn set_search_response(&self, response: SearchResponse) {
26        let mut guard = self.search_response.lock().unwrap();
27        *guard = Some(response);
28        // Reset search count when setting a new response
29        *self.search_count.lock().unwrap() = 0;
30    }
31
32    /// Clear the configured response.
33    pub fn clear_response(&self) {
34        let mut guard = self.search_response.lock().unwrap();
35        *guard = None;
36        // Reset search count when clearing
37        *self.search_count.lock().unwrap() = 0;
38    }
39}
40
41#[async_trait]
42impl Source for MockSource {
43    fn id(&self) -> &str {
44        "mock"
45    }
46
47    fn name(&self) -> &str {
48        "Mock Source"
49    }
50
51    fn capabilities(&self) -> SourceCapabilities {
52        SourceCapabilities::SEARCH
53    }
54
55    async fn search(&self, query: &SearchQuery) -> Result<SearchResponse, SourceError> {
56        // Increment search count
57        let mut count = self.search_count.lock().unwrap();
58        *count += 1;
59        let is_first_call = *count == 1;
60        drop(count);
61
62        let guard = self.search_response.lock().unwrap();
63        match &*guard {
64            Some(response) => {
65                if is_first_call {
66                    Ok(response.clone())
67                } else {
68                    // Return empty results for subsequent calls (simulating pagination end)
69                    Ok(SearchResponse::new(Vec::new(), "Mock Source", &query.query))
70                }
71            }
72            None => Ok(SearchResponse::new(Vec::new(), "Mock Source", &query.query)),
73        }
74    }
75}
76
77/// Helper function to create a mock paper for testing.
78pub fn make_paper(paper_id: &str, title: &str, source_type: SourceType) -> Paper {
79    Paper::new(
80        paper_id.to_string(),
81        title.to_string(),
82        format!("http://example.com/{}", paper_id),
83        source_type,
84    )
85}