research_master/sources/
mock.rs1use async_trait::async_trait;
4use std::sync::Mutex;
5
6use crate::models::{Paper, SearchQuery, SearchResponse, SourceType};
7use crate::sources::{Source, SourceCapabilities, SourceError};
8
9#[derive(Debug, Default, Clone)]
11pub struct MockSource {
12 search_response: std::sync::Arc<Mutex<Option<SearchResponse>>>,
15 search_count: std::sync::Arc<Mutex<usize>>,
16}
17
18impl MockSource {
19 pub fn new() -> Self {
21 Self::default()
22 }
23
24 pub fn set_search_response(&self, response: SearchResponse) {
26 let mut guard = self.search_response.lock().unwrap();
27 *guard = Some(response);
28 *self.search_count.lock().unwrap() = 0;
30 }
31
32 pub fn clear_response(&self) {
34 let mut guard = self.search_response.lock().unwrap();
35 *guard = None;
36 *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 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 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
77pub 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}