1use std::path::PathBuf;
16use std::sync::Arc;
17
18use clap::{Parser, ValueEnum};
19use sqlness::interceptor::Registry;
20use sqlness::{ConfigBuilder, Runner};
21
22use crate::cmd::SqlnessConfig;
23use crate::env::bare::{Env, ServiceProvider, StoreConfig, WalConfig};
24use crate::{protocol_interceptor, util};
25
26#[derive(ValueEnum, Debug, Clone)]
27#[clap(rename_all = "snake_case")]
28enum Wal {
29 RaftEngine,
30 Kafka,
31}
32
33#[derive(clap::Args, Debug, Clone, Default)]
35pub(crate) struct ServerAddr {
36 #[clap(short, long)]
38 pub(crate) server_addr: Option<String>,
39
40 #[clap(short, long, requires = "server_addr")]
42 pub(crate) pg_server_addr: Option<String>,
43
44 #[clap(short, long, requires = "server_addr")]
46 pub(crate) mysql_server_addr: Option<String>,
47}
48
49#[derive(Debug, Parser)]
50pub struct BareCommand {
52 #[clap(flatten)]
53 config: SqlnessConfig,
54
55 #[command(flatten)]
57 server_addr: ServerAddr,
58
59 #[clap(short, long, default_value = "raft_engine")]
61 wal: Wal,
62
63 #[clap(short, long)]
66 kafka_wal_broker_endpoints: Option<String>,
67
68 #[clap(long)]
71 bins_dir: Option<PathBuf>,
72
73 #[clap(long)]
76 preserve_state: bool,
77
78 #[clap(long, default_value = "true")]
80 pull_version_on_need: bool,
81
82 #[clap(long)]
84 store_addrs: Vec<String>,
85
86 #[clap(long, default_value = "false")]
88 setup_etcd: bool,
89
90 #[clap(long, default_missing_value = "", num_args(0..=1))]
92 setup_pg: Option<ServiceProvider>,
93
94 #[clap(long, default_missing_value = "", num_args(0..=1))]
96 setup_mysql: Option<ServiceProvider>,
97
98 #[clap(short, long, default_value = "0")]
100 jobs: usize,
101
102 #[clap(long)]
104 extra_args: Vec<String>,
105
106 #[clap(long, default_value = "false")]
108 enable_flat_format: bool,
109
110 #[clap(long, default_value = "false")]
112 enable_gc: bool,
113}
114
115#[derive(Debug, PartialEq, Eq)]
116struct Parallelism {
117 env: usize,
119 instance: usize,
121}
122
123impl BareCommand {
124 fn parallelism(&self) -> Parallelism {
125 let self_managed_kafka =
126 matches!(self.wal, Wal::Kafka) && self.kafka_wal_broker_endpoints.is_none();
127 if self.server_addr.server_addr.is_some()
131 || self.config.test_filter != ".*"
132 || self_managed_kafka
133 {
134 return Parallelism {
135 env: 1,
136 instance: 1,
137 };
138 }
139
140 if self.setup_etcd
145 || !self.store_addrs.is_empty()
146 || self.setup_pg.is_some()
147 || self.setup_mysql.is_some()
148 || self.kafka_wal_broker_endpoints.is_some()
149 {
150 return Parallelism {
151 env: self.jobs,
152 instance: 1,
153 };
154 }
155
156 Parallelism {
157 env: self.jobs,
158 instance: self.jobs,
159 }
160 }
161
162 pub async fn run(mut self) {
163 let temp_dir = tempfile::Builder::new()
164 .prefix("sqlness")
165 .tempdir()
166 .unwrap();
167 let sqlness_home = temp_dir.keep();
168
169 let mut interceptor_registry: Registry = Default::default();
170 interceptor_registry.register(
171 protocol_interceptor::PREFIX,
172 Arc::new(protocol_interceptor::ProtocolInterceptorFactory),
173 );
174
175 if let Some(d) = &self.config.case_dir
176 && !d.is_dir()
177 {
178 panic!("{} is not a directory", d.display());
179 }
180 if self.jobs == 0 {
181 self.jobs = num_cpus::get() / 2;
182 }
183 let parallelism = self.parallelism();
184 println!(
185 "Environment parallelism: {}, instance parallelism per environment: {}",
186 parallelism.env, parallelism.instance
187 );
188
189 let config = ConfigBuilder::default()
190 .case_dir(util::get_case_dir(self.config.case_dir))
191 .fail_fast(self.config.fail_fast)
192 .test_filter(self.config.test_filter)
193 .follow_links(true)
194 .env_config_file(self.config.env_config_file)
195 .interceptor_registry(interceptor_registry)
196 .parallelism(parallelism.instance)
197 .env_parallelism(parallelism.env)
198 .build()
199 .unwrap();
200
201 let wal = match self.wal {
202 Wal::RaftEngine => WalConfig::RaftEngine,
203 Wal::Kafka => WalConfig::Kafka {
204 needs_kafka_cluster: self.kafka_wal_broker_endpoints.is_none(),
205 broker_endpoints: self
206 .kafka_wal_broker_endpoints
207 .map(|s| s.split(',').map(|s| s.to_string()).collect())
208 .unwrap_or(vec!["127.0.0.1:9092".to_string()]),
210 },
211 };
212
213 let store = StoreConfig {
214 store_addrs: self.store_addrs.clone(),
215 setup_etcd: self.setup_etcd,
216 setup_pg: self.setup_pg,
217 setup_mysql: self.setup_mysql,
218 enable_flat_format: self.enable_flat_format,
219 enable_gc: self.enable_gc,
220 };
221
222 let runner = Runner::new(
223 config,
224 Env::new(
225 sqlness_home.clone(),
226 self.server_addr,
227 wal,
228 self.pull_version_on_need,
229 self.bins_dir,
230 store,
231 self.extra_args,
232 ),
233 );
234 match runner.run().await {
235 Ok(_) => println!("\x1b[32mAll sqlness tests passed!\x1b[0m"),
236 Err(e) => {
237 println!("\x1b[31mTest failed: {}\x1b[0m", e);
238 std::process::exit(1);
239 }
240 }
241
242 if !self.preserve_state {
244 if self.setup_etcd {
245 println!("Stopping etcd");
246 util::stop_rm_etcd();
247 }
248 println!("Removing state in {:?}", sqlness_home);
250 tokio::fs::remove_dir_all(sqlness_home).await.unwrap();
251 }
252 }
253}
254
255#[cfg(test)]
256mod tests {
257 use super::*;
258
259 fn parallelism(jobs: &str, args: &[&str]) -> Parallelism {
260 let cmd = BareCommand::try_parse_from(["bare", "-j", jobs].iter().chain(args)).unwrap();
261 cmd.parallelism()
262 }
263
264 #[test]
265 fn test_parallelism() {
266 let serial = Parallelism {
267 env: 1,
268 instance: 1,
269 };
270 let env_only = Parallelism {
271 env: 4,
272 instance: 1,
273 };
274
275 assert_eq!(
276 parallelism("4", &[]),
277 Parallelism {
278 env: 4,
279 instance: 4
280 }
281 );
282 assert_eq!(parallelism("4", &["--setup-pg"]), env_only);
283 assert_eq!(parallelism("4", &["--setup-mysql"]), env_only);
284 assert_eq!(parallelism("4", &["--setup-etcd"]), env_only);
285 assert_eq!(
286 parallelism("4", &["--store-addrs", "127.0.0.1:2379"]),
287 env_only
288 );
289 assert_eq!(
290 parallelism("4", &["-w", "kafka", "-k", "127.0.0.1:9092"]),
291 env_only
292 );
293 assert_eq!(parallelism("4", &["-w", "kafka"]), serial);
294 assert_eq!(parallelism("4", &["-t", "basic"]), serial);
295 assert_eq!(
296 parallelism(
297 "4",
298 &[
299 "-s",
300 "127.0.0.1:4001",
301 "-p",
302 "127.0.0.1:4003",
303 "-m",
304 "127.0.0.1:4002"
305 ]
306 ),
307 serial
308 );
309 assert_eq!(parallelism("1", &["--setup-pg"]), serial);
310 }
311}