1use std::ffi::{OsStr, OsString};
2use std::path::{Path, PathBuf};
3use std::time::Duration;
4
5use crate::error::ErrorInner;
6use crate::util::{KillOnDrop, TempDir};
7use crate::Error;
8use std::process::Command;
9
10pub struct TempMongo {
15 tempdir: TempDir,
16 socket_path: PathBuf,
17 log_path: PathBuf,
18 client: mongodb::Client,
19 server: KillOnDrop,
20}
21
22impl std::fmt::Debug for TempMongo {
23 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24 f.debug_struct("TempMongo")
25 .field("tempdir", &self.tempdir.path())
26 .field("socket_path", &self.socket_path())
27 .field("log_path", &self.log_path())
28 .field("server_pid", &self.server.id())
29 .finish_non_exhaustive()
30 }
31}
32
33impl TempMongo {
34 pub async fn new() -> Result<Self, Error> {
36 Self::from_builder(&TempMongoBuilder::new()).await
37 }
38
39 pub fn builder() -> TempMongoBuilder {
43 TempMongoBuilder::new()
44 }
45
46 pub fn process_id(&self) -> u32 {
48 self.server.id()
49 }
50
51 pub fn directory(&self) -> &Path {
53 self.tempdir.path()
54 }
55
56 pub fn socket_path(&self) -> &Path {
58 &self.socket_path
59 }
60
61 pub fn log_path(&self) -> &Path {
63 &self.log_path
64 }
65
66 pub fn client(&self) -> &mongodb::Client {
71 &self.client
72 }
73
74 pub fn set_clean_on_drop(&mut self, clean_on_drop: bool) {
76 self.tempdir.set_clean_on_drop(clean_on_drop);
77 }
78
79 pub async fn kill_and_clean(mut self) -> Result<(), Error> {
87 self.client.shutdown_immediate().await;
88 self.server.kill()
89 .map_err(ErrorInner::KillServer)?;
90
91 let path = self.tempdir.path().to_owned();
92 self.tempdir.close()
93 .map_err(|e| ErrorInner::CleanDir(path, e))?;
94 Ok(())
95 }
96
97 pub async fn kill_no_clean(mut self) -> Result<(), Error> {
104 let _path = self.tempdir.into_path();
105 self.client.shutdown_immediate().await;
106 self.server.kill()
107 .map_err(ErrorInner::KillServer)?;
108 Ok(())
109 }
110
111 async fn from_builder(builder: &TempMongoBuilder) -> Result<Self, Error> {
113 let tempdir = builder.make_temp_dir().map_err(ErrorInner::MakeTempDir)?;
114 let db_dir = tempdir.path().join("db");
115 let socket_path = tempdir.path().join("mongod.sock");
116 let log_path = tempdir.path().join("mongod.log");
117
118 std::fs::create_dir(&db_dir)
119 .map_err(|e| ErrorInner::MakeDbDir(db_dir.clone(), e))?;
120
121 let server = Command::new(builder.get_command())
122 .arg("--bind_ip")
123 .arg(&socket_path)
124 .arg("--dbpath")
125 .arg(db_dir)
126 .arg("--logpath")
127 .arg(&log_path)
128 .arg("--nounixsocket")
129 .arg("--noauth")
130 .spawn()
131 .map_err(|e| ErrorInner::SpawnServer(builder.get_command_string(), e))?;
132 let server = KillOnDrop::new(server);
133
134 let client_options = mongodb::options::ClientOptions::builder()
135 .hosts(vec![mongodb::options::ServerAddress::Unix { path: socket_path.clone() }])
136 .connect_timeout(Duration::from_millis(10))
137 .retry_reads(false)
138 .retry_writes(false)
139 .build();
140 let client = mongodb::Client::with_options(client_options)
141 .map_err(|e| ErrorInner::Connect(socket_path.display().to_string(), e))?;
142
143 client.list_databases(None, None).await
144 .map_err(|e| ErrorInner::Connect(socket_path.display().to_string(), e))?;
145
146 Ok(Self {
147 tempdir,
148 socket_path,
149 log_path,
150 server,
151 client,
152 })
153 }
154}
155
156#[derive(Debug)]
160pub struct TempMongoBuilder {
161 parent_directory: Option<PathBuf>,
165
166 clean_on_drop: bool,
168
169 command: Option<OsString>,
171}
172
173impl TempMongoBuilder {
174 pub fn new() -> Self {
176 Self {
177 parent_directory: None,
178 command: None,
179 clean_on_drop: true,
180 }
181 }
182
183 pub async fn spawn(&self) -> Result<TempMongo, Error> {
185 TempMongo::from_builder(self).await
186 }
187
188 pub fn clean_on_drop(mut self, clean_on_drop: bool) -> Self {
192 self.clean_on_drop = clean_on_drop;
193 self
194 }
195
196 pub fn mongod_command(mut self, command: impl Into<OsString>) -> Self {
200 self.command = Some(command.into());
201 self
202 }
203
204 fn get_command(&self) -> &OsStr {
206 self.command
207 .as_deref()
208 .unwrap_or("mongod".as_ref())
209 }
210
211 fn get_command_string(&self) -> String {
213 self.get_command().to_string_lossy().into()
214 }
215
216 fn make_temp_dir(&self) -> std::io::Result<TempDir> {
218 match &self.parent_directory {
219 Some(dir) => TempDir::new_in(dir, self.clean_on_drop),
220 None => TempDir::new(self.clean_on_drop),
221 }
222 }
223}
224
225impl Default for TempMongoBuilder {
226 fn default() -> Self {
227 Self::new()
228 }
229}