Skip to content

Commit 38542b9

Browse files
TanmayArya-1pTanmayArya-1p
authored
fix: extend lifetime of tempdir to prevent deletion (#62)
1 parent 97a8a41 commit 38542b9

6 files changed

Lines changed: 32 additions & 53 deletions

File tree

‎Cargo.lock‎

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/api/src/lib.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -168,23 +168,23 @@ mod tests {
168168

169169
use super::*;
170170
use defs::ContentType;
171-
use tempfile::tempdir;
171+
use tempfile::{TempDir, tempdir};
172172

173173
// Helper function to create a test database
174-
fn create_test_db() -> VectorDb {
174+
fn create_test_db() -> (VectorDb, TempDir) {
175175
let temp_dir = tempdir().unwrap();
176176
let config = DbConfig {
177177
storage_type: StorageType::RocksDb,
178178
index_type: IndexType::Flat,
179179
data_path: temp_dir.path().to_path_buf(),
180180
dimension: 3,
181181
};
182-
init_api(config).unwrap()
182+
(init_api(config).unwrap(), temp_dir)
183183
}
184184

185185
#[test]
186186
fn test_insert_and_get() {
187-
let db = create_test_db();
187+
let (db, _temp_dir) = create_test_db();
188188
let vector = vec![1.0, 2.0, 3.0];
189189
let payload = Payload {
190190
content_type: ContentType::Text,
@@ -209,7 +209,7 @@ mod tests {
209209

210210
#[test]
211211
fn test_dimension_mismatch() {
212-
let db = create_test_db();
212+
let (db, _temp_dir) = create_test_db();
213213
let v1 = vec![1.0, 2.0, 3.0];
214214
let v2 = vec![1.0, 2.0];
215215
let payload = defs::Payload {
@@ -228,7 +228,7 @@ mod tests {
228228

229229
#[test]
230230
fn test_delete() {
231-
let db = create_test_db();
231+
let (db, _temp_dir) = create_test_db();
232232
let vector = vec![1.0, 2.0, 3.0];
233233
let payload = Payload {
234234
content_type: ContentType::Text,
@@ -251,7 +251,7 @@ mod tests {
251251

252252
#[test]
253253
fn test_search() {
254-
let db = create_test_db();
254+
let (db, _temp_dir) = create_test_db();
255255

256256
// Insert some points
257257
let vectors = vec![
@@ -280,7 +280,7 @@ mod tests {
280280

281281
#[test]
282282
fn test_search_limit() {
283-
let db = create_test_db();
283+
let (db, _temp_dir) = create_test_db();
284284

285285
// Insert 5 points
286286
let mut ids = Vec::new();
@@ -307,7 +307,7 @@ mod tests {
307307

308308
#[test]
309309
fn test_empty_database() {
310-
let db = create_test_db();
310+
let (db, _temp_dir) = create_test_db();
311311

312312
// Get non-existent point
313313
assert!(db.get(Uuid::new_v4()).unwrap().is_none());
@@ -319,7 +319,7 @@ mod tests {
319319

320320
#[test]
321321
fn test_list_vectors() {
322-
let db = create_test_db();
322+
let (db, _temp_dir) = create_test_db();
323323
// insert some points
324324
let mut ids = Vec::new();
325325
for i in 0..10 {
@@ -350,7 +350,7 @@ mod tests {
350350

351351
#[test]
352352
fn test_build_index() {
353-
let db = create_test_db();
353+
let (db, _temp_dir) = create_test_db();
354354

355355
// insert some points
356356
for i in 0..10 {

‎crates/grpc/src/tests.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use index::IndexType;
88
use std::net::SocketAddr;
99
use std::sync::Arc;
1010
use storage::StorageType;
11-
use tempfile::tempdir;
11+
use tempfile::{TempDir, tempdir};
1212
use tonic::transport::Channel;
1313

1414
// Inspired from https://github.com/hyperium/tonic/discussions/924#discussioncomment-9854088
@@ -22,7 +22,7 @@ fn append_test_auth_header<T>(request: &mut tonic::Request<T>, token: &str) {
2222
.insert(AUTHORIZATION_HEADER_KEY, auth_value.parse().unwrap());
2323
}
2424

25-
async fn start_test_server() -> Result<SocketAddr, Box<dyn std::error::Error>> {
25+
async fn start_test_server() -> Result<(SocketAddr, TempDir), Box<dyn std::error::Error>> {
2626
// using a temporary directory for db datapath
2727
let temp_dir = tempdir().unwrap();
2828

@@ -50,7 +50,7 @@ async fn start_test_server() -> Result<SocketAddr, Box<dyn std::error::Error>> {
5050
.inspect_err(|err| panic!("Could not start test server : {:?}", err));
5151
});
5252

53-
Ok(listener_addr)
53+
Ok((listener_addr, temp_dir))
5454
}
5555

5656
async fn create_test_client(
@@ -65,7 +65,7 @@ async fn create_test_client(
6565

6666
#[tokio::test]
6767
async fn test_grpc_server_start() {
68-
let server_addr = start_test_server().await.unwrap();
68+
let (server_addr, _temp_dir) = start_test_server().await.unwrap();
6969
let mut client = create_test_client(server_addr).await.unwrap();
7070

7171
// insert a test vector
@@ -85,7 +85,7 @@ async fn test_grpc_server_start() {
8585

8686
#[tokio::test]
8787
async fn test_insert_vector_rpc() {
88-
let server_addr = start_test_server().await.unwrap();
88+
let (server_addr, _temp_dir) = start_test_server().await.unwrap();
8989
let mut client = create_test_client(server_addr).await.unwrap();
9090

9191
// insert a test vector
@@ -133,7 +133,7 @@ async fn test_insert_vector_rpc() {
133133

134134
#[tokio::test]
135135
async fn test_delete_vector_rpc() {
136-
let server_addr = start_test_server().await.unwrap();
136+
let (server_addr, _temp_dir) = start_test_server().await.unwrap();
137137
let mut client = create_test_client(server_addr).await.unwrap();
138138

139139
// insert a test vector
@@ -175,7 +175,7 @@ async fn test_delete_vector_rpc() {
175175

176176
#[tokio::test]
177177
async fn test_search_vector_rpc() {
178-
let server_addr = start_test_server().await.unwrap();
178+
let (server_addr, _temp_dir) = start_test_server().await.unwrap();
179179
let mut client = create_test_client(server_addr).await.unwrap();
180180

181181
// insert a test vector
@@ -221,7 +221,7 @@ async fn test_search_vector_rpc() {
221221

222222
#[tokio::test]
223223
async fn test_unauthorized_rpc() {
224-
let server_addr = start_test_server().await.unwrap();
224+
let (server_addr, _temp_dir) = start_test_server().await.unwrap();
225225
let mut client = create_test_client(server_addr).await.unwrap();
226226

227227
// insert a test vector

‎crates/server/Cargo.toml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ prost.workspace = true
1818
serde.workspace = true
1919
serde_json.workspace = true
2020
storage.workspace = true
21-
tempfile.workspace = true
2221
tokio.workspace = true
2322
tokio-stream.workspace = true
2423
tonic.workspace = true

‎crates/server/src/config.rs‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ use std::fs;
66
use std::net::SocketAddr;
77
use std::path::PathBuf;
88
use storage::StorageType;
9-
use tempfile::tempdir;
109
use tracing::{Level, event};
1110

1211
const DEFAULT_HTTP_PORT: &str = "3000";
@@ -146,11 +145,7 @@ impl ServerConfig {
146145
fs::create_dir_all(&path).map_err(|_| ConfigError::InvalidDataPath)?;
147146
path
148147
} else {
149-
let tempbuf = tempdir()
150-
.map_err(|e| ConfigError::IoError(e.to_string()))?
151-
.path()
152-
.to_path_buf()
153-
.join("vectordb");
148+
let tempbuf = env::temp_dir().join("vectordb");
154149
fs::create_dir_all(&tempbuf).map_err(|e| ConfigError::IoError(e.to_string()))?;
155150
event!(
156151
Level::WARN,

‎crates/storage/src/rocks_db.rs‎

Lines changed: 12 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -160,26 +160,24 @@ mod tests {
160160
use defs::ContentType;
161161
use uuid::Uuid;
162162

163-
use tempfile::tempdir;
163+
use tempfile::{TempDir, tempdir};
164164

165-
fn create_test_db() -> (RocksDbStorage, String) {
165+
fn create_test_db() -> (RocksDbStorage, TempDir) {
166166
let temp_dir = tempdir().unwrap();
167-
let temp_dir_path = temp_dir.path().to_str().unwrap().to_string();
168167

169-
let db = RocksDbStorage::new(temp_dir_path.clone()).expect("Failed to create RocksDB");
170-
(db, temp_dir_path)
168+
let db = RocksDbStorage::new(temp_dir.path()).expect("Failed to create RocksDB");
169+
(db, temp_dir)
171170
}
172171

173172
#[test]
174173
fn test_new_rocksdb_storage() {
175-
let (db, path) = create_test_db();
176-
assert_eq!(db.get_current_path(), PathBuf::from(path.clone()));
177-
std::fs::remove_dir_all(path).unwrap_or_default();
174+
let (db, temp_dir) = create_test_db();
175+
assert_eq!(db.get_current_path(), PathBuf::from(temp_dir.path()));
178176
}
179177

180178
#[test]
181179
fn test_insert_and_get_vector() {
182-
let (db, path) = create_test_db();
180+
let (db, _temp_dir) = create_test_db();
183181
let id = Uuid::new_v4();
184182
let vector = Some(vec![0.1, 0.2, 0.3]);
185183
let payload = Some(Payload {
@@ -190,13 +188,11 @@ mod tests {
190188
assert!(db.insert_point(id, vector.clone(), payload).is_ok());
191189
let result = db.get_vector(id).unwrap();
192190
assert_eq!(result, vector);
193-
194-
std::fs::remove_dir_all(path).unwrap_or_default();
195191
}
196192

197193
#[test]
198194
fn test_insert_and_get_payload() {
199-
let (db, path) = create_test_db();
195+
let (db, _temp_dir) = create_test_db();
200196
let id = Uuid::new_v4();
201197
let payload = Some(Payload {
202198
content_type: ContentType::Text,
@@ -212,13 +208,11 @@ mod tests {
212208
content: "Test".to_string(),
213209
});
214210
assert_eq!(result, expected);
215-
216-
std::fs::remove_dir_all(path).unwrap_or_default();
217211
}
218212

219213
#[test]
220214
fn test_contains_point() {
221-
let (db, path) = create_test_db();
215+
let (db, _temp_dir) = create_test_db();
222216
let id = Uuid::new_v4();
223217
let payload = Some(Payload {
224218
content_type: ContentType::Text,
@@ -231,13 +225,11 @@ mod tests {
231225
db.insert_point(id, vector, payload).unwrap();
232226

233227
assert!(db.contains_point(id).unwrap());
234-
235-
std::fs::remove_dir_all(path).unwrap_or_default();
236228
}
237229

238230
#[test]
239231
fn test_delete_point() {
240-
let (db, path) = create_test_db();
232+
let (db, _temp_dir) = create_test_db();
241233
let id = Uuid::new_v4();
242234
let payload = Some(Payload {
243235
content_type: ContentType::Text,
@@ -255,27 +247,21 @@ mod tests {
255247
assert!(!db.contains_point(id).unwrap());
256248
assert_eq!(db.get_vector(id).unwrap(), None);
257249
assert_eq!(db.get_payload(id).unwrap(), None);
258-
259-
std::fs::remove_dir_all(path).unwrap_or_default();
260250
}
261251

262252
#[test]
263253
fn test_get_nonexistent_vector() {
264-
let (db, path) = create_test_db();
254+
let (db, _temp_dir) = create_test_db();
265255
let id = Uuid::new_v4();
266256

267257
assert_eq!(db.get_vector(id).unwrap(), None);
268-
269-
std::fs::remove_dir_all(path).unwrap_or_default();
270258
}
271259

272260
#[test]
273261
fn test_get_nonexistent_payload() {
274-
let (db, path) = create_test_db();
262+
let (db, _temp_dir) = create_test_db();
275263
let id = Uuid::new_v4();
276264

277265
assert_eq!(db.get_payload(id).unwrap(), None);
278-
279-
std::fs::remove_dir_all(path).unwrap_or_default();
280266
}
281267
}

0 commit comments

Comments
 (0)