1use std::collections::BTreeMap;
7
8#[derive(Debug, Clone, PartialEq)]
17pub struct MockTable {
18 pub name: String,
20 pub rows: Vec<BTreeMap<String, String>>,
22}
23
24impl MockTable {
25 pub fn new(name: impl Into<String>) -> Self {
27 Self {
28 name: name.into(),
29 rows: Vec::new(),
30 }
31 }
32
33 pub fn add_row(&mut self, row: BTreeMap<String, String>) {
35 self.rows.push(row);
36 }
37
38 pub fn to_sql(&self) -> String {
43 if self.rows.is_empty() {
44 return format!("SELECT NULL WHERE 1=0 /* empty mock: {} */", self.name);
45 }
46
47 let columns: Vec<&String> = self.rows[0].keys().collect();
49
50 let mut parts = Vec::new();
51
52 for (i, row) in self.rows.iter().enumerate() {
53 let values: Vec<String> = columns
54 .iter()
55 .map(|col| {
56 let val = row.get(*col).map(|s| s.as_str()).unwrap_or("NULL");
57 if val == "NULL" {
58 if i == 0 {
59 format!("NULL AS {col}")
60 } else {
61 "NULL".to_string()
62 }
63 } else if i == 0 {
64 format!("'{}' AS {col}", val.replace('\'', "''"))
65 } else {
66 format!("'{}'", val.replace('\'', "''"))
67 }
68 })
69 .collect();
70 parts.push(format!("SELECT {}", values.join(", ")));
71 }
72
73 parts.join("\nUNION ALL\n")
74 }
75}
76
77#[macro_export]
88macro_rules! mock_rows {
89 [$( {$($key:literal => $val:literal),* $(,)?} ),* $(,)?] => {
90 vec![
91 $(
92 {
93 let mut row = ::std::collections::BTreeMap::new();
94 $( row.insert($key.to_string(), $val.to_string()); )*
95 row
96 }
97 ),*
98 ]
99 };
100}
101
102#[cfg(test)]
103mod tests {
104 use super::*;
105
106 #[test]
107 fn test_mock_table_single_row() {
108 let mut mock = MockTable::new("users");
109 let mut row = BTreeMap::new();
110 row.insert("id".to_string(), "1".to_string());
111 row.insert("name".to_string(), "Alice".to_string());
112 mock.add_row(row);
113
114 let sql = mock.to_sql();
115 assert_eq!(sql, "SELECT '1' AS id, 'Alice' AS name");
116 }
117
118 #[test]
119 fn test_mock_table_multiple_rows() {
120 let mut mock = MockTable::new("users");
121
122 let mut row1 = BTreeMap::new();
123 row1.insert("id".to_string(), "1".to_string());
124 row1.insert("name".to_string(), "Alice".to_string());
125 mock.add_row(row1);
126
127 let mut row2 = BTreeMap::new();
128 row2.insert("id".to_string(), "2".to_string());
129 row2.insert("name".to_string(), "Bob".to_string());
130 mock.add_row(row2);
131
132 let sql = mock.to_sql();
133 assert_eq!(
134 sql,
135 "SELECT '1' AS id, 'Alice' AS name\nUNION ALL\nSELECT '2', 'Bob'"
136 );
137 }
138
139 #[test]
140 fn test_mock_table_with_null() {
141 let mut mock = MockTable::new("users");
142 let mut row = BTreeMap::new();
143 row.insert("id".to_string(), "1".to_string());
144 row.insert("email".to_string(), "NULL".to_string());
145 mock.add_row(row);
146
147 let sql = mock.to_sql();
148 assert!(sql.contains("NULL AS email"));
149 assert!(!sql.contains("'NULL'"));
150 }
151
152 #[test]
153 fn test_mock_table_empty() {
154 let mock = MockTable::new("users");
155 let sql = mock.to_sql();
156 assert!(sql.contains("WHERE 1=0"));
157 }
158
159 #[test]
160 fn test_mock_table_sql_injection_safe() {
161 let mut mock = MockTable::new("users");
162 let mut row = BTreeMap::new();
163 row.insert("name".to_string(), "O'Brien".to_string());
164 mock.add_row(row);
165
166 let sql = mock.to_sql();
167 assert!(sql.contains("O''Brien"));
168 }
169
170 #[test]
171 fn test_mock_rows_macro() {
172 let rows = mock_rows![
173 {"id" => "1", "name" => "Alice"},
174 {"id" => "2", "name" => "Bob"},
175 ];
176 assert_eq!(rows.len(), 2);
177 assert_eq!(rows[0].get("id").unwrap(), "1");
178 assert_eq!(rows[1].get("name").unwrap(), "Bob");
179 }
180}