on iko/fix help sign in
reshape/tests/common.rs
310 lines 9.2 kB history raw
1use colored::Colorize;
2use postgres::{Client, NoTls};
3use reshape::{migrations::Migration, Reshape};
5pub struct Test<'a> {
6 name: &'a str,
7 reshape: Reshape,
8 old_db: Client,
9 new_db: Client,
11 first_migration: Option<Migration>,
12 second_migration: Option<Migration>,
13 expect_failure: bool,
15 clear_fn: Option<fn(&mut Client) -> ()>,
16 after_first_fn: Option<fn(&mut Client) -> ()>,
17 intermediate_fn: Option<fn(&mut Client, &mut Client) -> ()>,
18 after_completion_fn: Option<fn(&mut Client) -> ()>,
19 after_abort_fn: Option<fn(&mut Client) -> ()>,
22impl Test<'_> {
23 pub fn new<'a>(name: &'a str) -> Test<'a> {
24 let connection_string = std::env::var("POSTGRES_CONNECTION_STRING")
25 .unwrap_or("postgres://postgres:postgres@localhost/reshape_test".to_string());
27 let old_db = Client::connect(&connection_string, NoTls).unwrap();
28 let new_db = Client::connect(&connection_string, NoTls).unwrap();
30 let reshape = Reshape::new(&connection_string).unwrap();
32 Test {
33 name,
34 reshape,
35 old_db,
36 new_db,
37 first_migration: None,
38 second_migration: None,
39 expect_failure: false,
40 clear_fn: None,
41 after_first_fn: None,
42 intermediate_fn: None,
43 after_completion_fn: None,
44 after_abort_fn: None,
45 }
46 }
48 pub fn first_migration(&mut self, migration: &str) -> &mut Self {
49 self.first_migration = Some(Self::parse_migration(migration));
50 self
51 }
53 #[allow(dead_code)]
54 pub fn second_migration(&mut self, migration: &str) -> &mut Self {
55 self.second_migration = Some(Self::parse_migration(migration));
56 self
57 }
59 #[allow(dead_code)]
60 pub fn clear(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
61 self.clear_fn = Some(f);
62 self
63 }
65 #[allow(dead_code)]
66 pub fn after_first(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
67 self.after_first_fn = Some(f);
68 self
69 }
71 #[allow(dead_code)]
72 pub fn intermediate(&mut self, f: fn(&mut Client, &mut Client) -> ()) -> &mut Self {
73 self.intermediate_fn = Some(f);
74 self
75 }
77 #[allow(dead_code)]
78 pub fn after_completion(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
79 self.after_completion_fn = Some(f);
80 self
81 }
83 #[allow(dead_code)]
84 pub fn after_abort(&mut self, f: fn(&mut Client) -> ()) -> &mut Self {
85 self.after_abort_fn = Some(f);
86 self
87 }
89 #[allow(dead_code)]
90 pub fn expect_failure(&mut self) {
91 self.expect_failure = true;
92 }
94 fn parse_migration(encoded: &str) -> Migration {
95 toml::from_str(encoded).unwrap()
96 }
99enum RunType {
100 Simple,
101 Completion,
102 Abort,
105impl Test<'_> {
106 #[allow(dead_code)]
107 pub fn run(&mut self) {
108 if self.second_migration.is_some() {
109 // Run to completion
110 print_heading(&format!("Test completion: {}", self.name));
111 self.run_internal(RunType::Completion);
113 // Run and abort
114 print_heading(&format!("Test abort: {}", self.name));
115 self.run_internal(RunType::Abort);
116 } else {
117 print_heading(&format!("Test: {}", self.name));
118 self.run_internal(RunType::Simple);
119 }
120 }
122 fn run_internal(&mut self, run_type: RunType) {
123 print_subheading("Clearing database");
124 self.reshape.remove().unwrap();
126 if let Some(clear_fn) = self.clear_fn {
127 clear_fn(&mut self.old_db);
128 }
130 // Apply first migration, will automatically complete
131 print_subheading("Applying first migration");
132 let first_migration = self
133 .first_migration
134 .as_ref()
135 .expect("no starting migration set");
136 self.reshape.migrate(vec![first_migration.clone()]).unwrap();
138 // Update search path
139 self.old_db
140 .simple_query(&reshape::schema_query_for_migration(&first_migration.name))
141 .unwrap();
143 // Automatically complete first migration
144 self.reshape.complete().unwrap();
146 // Run setup function
147 if let Some(after_first_fn) = self.after_first_fn {
148 print_subheading("Running setup and first checks");
149 after_first_fn(&mut self.old_db);
150 print_success();
151 }
153 // Apply second migration
154 if let Some(second_migration) = &self.second_migration {
155 if self.expect_failure {
156 print_subheading("Applying second migration (expecting failure)");
157 let result = self
158 .reshape
159 .migrate(vec![first_migration.clone(), second_migration.clone()]);
161 if result.is_ok() {
162 panic!("expected second migration to fail");
163 }
164 } else {
165 print_subheading("Applying second migration");
166 self.reshape
167 .migrate(vec![first_migration.clone(), second_migration.clone()])
168 .unwrap();
169 }
171 // Update search path
172 self.new_db
173 .simple_query(&reshape::schema_query_for_migration(&second_migration.name))
174 .unwrap();
176 if let Some(intermediate_fn) = self.intermediate_fn {
177 print_subheading("Running intermediate checks");
178 intermediate_fn(&mut self.old_db, &mut self.new_db);
179 print_success();
180 }
182 match run_type {
183 RunType::Completion => {
184 print_subheading("Completing");
185 self.reshape.complete().unwrap();
187 if let Some(after_completion_fn) = self.after_completion_fn {
188 print_subheading("Running post-completion checks");
189 after_completion_fn(&mut self.new_db);
190 print_success();
191 }
192 }
193 RunType::Abort => {
194 print_subheading("Aborting");
195 self.reshape.abort().unwrap();
197 if let Some(after_abort_fn) = self.after_abort_fn {
198 print_subheading("Running post-abort checks");
199 after_abort_fn(&mut self.old_db);
200 print_success();
201 }
202 }
203 _ => {}
204 }
205 }
207 print_subheading("Checking cleanup");
208 assert_cleaned_up(&mut self.new_db);
209 print_success();
210 }
213fn print_heading(text: &str) {
214 let delimiter = std::iter::repeat("=").take(80).collect::<String>();
216 println!();
217 println!();
218 println!("{}", delimiter.blue().bold());
219 println!("{}", add_spacer(text, "=").blue().bold());
220 println!("{}", delimiter.blue().bold());
223fn print_subheading(text: &str) {
224 println!();
225 println!("{}", add_spacer(text, "=").blue());
228fn print_success() {
229 println!("{}", add_spacer("Success", "=").green());
232fn add_spacer(text: &str, char: &str) -> String {
233 const TARGET_WIDTH: usize = 80;
234 let num_of_chars = (TARGET_WIDTH - text.len() - 2) / 2;
235 let spacer = std::iter::repeat(char)
236 .take(num_of_chars)
237 .collect::<String>();
239 let extra = if text.len() % 2 == 0 { "" } else { char };
241 format!("{spacer} {text} {spacer}{extra}", spacer = spacer)
244pub fn assert_cleaned_up(db: &mut Client) {
245 // Make sure no temporary columns remain
246 let temp_columns: Vec<String> = db
247 .query(
248 "
249 SELECT column_name
250 FROM information_schema.columns
251 WHERE table_schema = 'public'
252 AND column_name LIKE '__reshape%'
253 ",
254 &[],
255 )
256 .unwrap()
257 .iter()
258 .map(|row| row.get(0))
259 .collect();
261 assert!(
262 temp_columns.is_empty(),
263 "expected no temporary columns to exist, found: {}",
264 temp_columns.join(", ")
265 );
267 // Make sure no triggers remain
268 let triggers: Vec<String> = db
269 .query(
270 "
271 SELECT trigger_name
272 FROM information_schema.triggers
273 WHERE trigger_schema = 'public'
274 AND trigger_name LIKE '__reshape%'
275 ",
276 &[],
277 )
278 .unwrap()
279 .iter()
280 .map(|row| row.get(0))
281 .collect();
283 assert!(
284 triggers.is_empty(),
285 "expected no triggers to exist, found: {}",
286 triggers.join(", ")
287 );
289 // Make sure no functions remain
290 let functions: Vec<String> = db
291 .query(
292 "
293 SELECT routine_name
294 FROM information_schema.routines
295 WHERE routine_schema = 'public'
296 AND routine_name LIKE '__reshape%'
297 ",
298 &[],
299 )
300 .unwrap()
301 .iter()
302 .map(|row| row.get(0))
303 .collect();
305 assert!(
306 functions.is_empty(),
307 "expected no functions to exist, found: {}",
308 functions.join(", ")
309 );