Skip to content

Commit 51cc1fd

Browse files
Arshdeep544adex
Arshdeep54
authored and
4adex
committed
refactor test.rs, error.rs, constants.rs for each crate (#89)
Signed-off-by: Arshdeep54 <balarsh535@gmail.com>
1 parent 7948327 commit 51cc1fd

34 files changed

Lines changed: 1154 additions & 1098 deletions

‎crates/api/src/lib.rs‎

Lines changed: 1 addition & 455 deletions
Large diffs are not rendered by default.

‎crates/api/src/tests.rs‎

Lines changed: 451 additions & 0 deletions
Large diffs are not rendered by default.

‎crates/http/src/auth.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,7 @@ use defs::ApiKeyRole;
99
use serde_json::json;
1010

1111
use crate::AppState;
12-
13-
const API_KEY_HEADER: &str = "api-key";
12+
use crate::constants::API_KEY_HEADER;
1413

1514
fn extract_key(req: &Request) -> Option<&str> {
1615
req.headers().get(API_KEY_HEADER)?.to_str().ok()

‎crates/http/src/constants.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
pub const API_KEY_HEADER: &str = "api-key";

‎crates/http/src/lib.rs‎

Lines changed: 2 additions & 173 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
pub mod auth;
2+
pub mod constants;
23
pub mod handler;
34

45
use api::VectorDb;
@@ -68,176 +69,4 @@ pub async fn run_http_server(
6869
}
6970

7071
#[cfg(test)]
71-
mod tests {
72-
use super::*;
73-
use api::DbConfig;
74-
use axum::http::StatusCode;
75-
use axum_test::TestServer;
76-
use defs::Similarity;
77-
use index::{IndexType, hnsw::HnswConfig, kd_tree::KDTreeConfig};
78-
use serde_json::json;
79-
use storage::StorageType;
80-
81-
use defs::{ApiKeyEntry, ApiKeyRole, ApiKeyStore};
82-
83-
const API_KEY: &str = "full-access-key";
84-
const READONLY_API_KEY: &str = "read-only-key";
85-
86-
fn test_db() -> Arc<api::VectorDb> {
87-
let temp_dir = tempfile::tempdir().unwrap();
88-
let db = api::init_api(DbConfig {
89-
storage_type: StorageType::InMemory,
90-
index_type: IndexType::Flat,
91-
data_path: temp_dir.path().to_path_buf(),
92-
dimension: 3,
93-
similarity: Similarity::Cosine,
94-
hnsw_config: HnswConfig::default(),
95-
kd_tree_config: KDTreeConfig::default(),
96-
})
97-
.unwrap();
98-
Arc::new(db)
99-
}
100-
101-
fn test_server() -> TestServer {
102-
let keys = ApiKeyStore::new(vec![
103-
ApiKeyEntry {
104-
name: "full".to_string(),
105-
role: ApiKeyRole::ReadWrite,
106-
key: API_KEY.to_string(),
107-
},
108-
ApiKeyEntry {
109-
name: "readonly".to_string(),
110-
role: ApiKeyRole::ReadOnly,
111-
key: READONLY_API_KEY.to_string(),
112-
},
113-
]);
114-
let router = create_router(test_db(), Arc::new(keys));
115-
TestServer::new(router).unwrap()
116-
}
117-
118-
#[tokio::test]
119-
async fn in_memory_storage_http_smoke_test() {
120-
let server = test_server();
121-
122-
let insert_response = server
123-
.post("/points")
124-
.add_header("api-key", API_KEY)
125-
.json(&json!({
126-
"vector": [1.0, 0.0, 0.0],
127-
"payload": {
128-
"content_type": "Text",
129-
"content": "smoke-test"
130-
}
131-
}))
132-
.await;
133-
insert_response.assert_status(StatusCode::CREATED);
134-
let insert_body: serde_json::Value = insert_response.json();
135-
let point_id = insert_body["point_id"].as_str().unwrap();
136-
137-
let get_response = server
138-
.get(&format!("/points/{point_id}"))
139-
.add_header("api-key", API_KEY)
140-
.await;
141-
get_response.assert_status_ok();
142-
let point_body: serde_json::Value = get_response.json();
143-
assert_eq!(point_body["payload"]["content"], "smoke-test");
144-
assert_eq!(point_body["vector"], json!([1.0, 0.0, 0.0]));
145-
146-
let search_response = server
147-
.post("/points/search")
148-
.add_header("api-key", API_KEY)
149-
.json(&json!({
150-
"vector": [1.0, 0.0, 0.0],
151-
"similarity": "Cosine",
152-
"limit": 1
153-
}))
154-
.await;
155-
search_response.assert_status_ok();
156-
let search_body: serde_json::Value = search_response.json();
157-
assert_eq!(search_body["results"], json!([point_id]));
158-
159-
let delete_response = server
160-
.delete(&format!("/points/{point_id}"))
161-
.add_header("api-key", API_KEY)
162-
.await;
163-
delete_response.assert_status(StatusCode::NO_CONTENT);
164-
165-
let missing_response = server
166-
.get(&format!("/points/{point_id}"))
167-
.add_header("api-key", API_KEY)
168-
.await;
169-
missing_response.assert_status(StatusCode::NOT_FOUND);
170-
}
171-
172-
#[tokio::test]
173-
async fn public_routes_require_no_key() {
174-
let server = test_server();
175-
server.get("/").await.assert_status_ok();
176-
server.get("/health").await.assert_status_ok();
177-
}
178-
179-
#[tokio::test]
180-
async fn protected_routes_reject_missing_key() {
181-
let server = test_server();
182-
server
183-
.get("/points/00000000-0000-0000-0000-000000000000")
184-
.await
185-
.assert_status(StatusCode::UNAUTHORIZED);
186-
server
187-
.post("/points")
188-
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}}))
189-
.await
190-
.assert_status(StatusCode::UNAUTHORIZED);
191-
}
192-
193-
#[tokio::test]
194-
async fn protected_routes_reject_wrong_key() {
195-
let server = test_server();
196-
server
197-
.get("/points/00000000-0000-0000-0000-000000000000")
198-
.add_header("api-key", "not-the-right-key")
199-
.await
200-
.assert_status(StatusCode::UNAUTHORIZED);
201-
}
202-
203-
#[tokio::test]
204-
async fn readonly_key_can_read_but_not_write() {
205-
let server = test_server();
206-
207-
let insert_response = server
208-
.post("/points")
209-
.add_header("api-key", API_KEY)
210-
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}}))
211-
.await;
212-
let point_id = insert_response.json::<serde_json::Value>()["point_id"]
213-
.as_str()
214-
.unwrap()
215-
.to_string();
216-
217-
server
218-
.get(&format!("/points/{point_id}"))
219-
.add_header("api-key", READONLY_API_KEY)
220-
.await
221-
.assert_status_ok();
222-
223-
server
224-
.post("/points/search")
225-
.add_header("api-key", READONLY_API_KEY)
226-
.json(&json!({"vector": [1.0, 0.0, 0.0], "similarity": "Cosine", "limit": 1}))
227-
.await
228-
.assert_status_ok();
229-
230-
server
231-
.delete(&format!("/points/{point_id}"))
232-
.add_header("api-key", READONLY_API_KEY)
233-
.await
234-
.assert_status(StatusCode::FORBIDDEN);
235-
236-
server
237-
.post("/points")
238-
.add_header("api-key", READONLY_API_KEY)
239-
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "y"}}))
240-
.await
241-
.assert_status(StatusCode::FORBIDDEN);
242-
}
243-
}
72+
mod tests;

‎crates/http/src/tests.rs‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
use super::*;
2+
use api::DbConfig;
3+
use axum::http::StatusCode;
4+
use axum_test::TestServer;
5+
use defs::Similarity;
6+
use index::{IndexType, hnsw::HnswConfig, kd_tree::KDTreeConfig};
7+
use serde_json::json;
8+
use storage::StorageType;
9+
10+
use defs::{ApiKeyEntry, ApiKeyRole, ApiKeyStore};
11+
12+
const API_KEY: &str = "full-access-key";
13+
const READONLY_API_KEY: &str = "read-only-key";
14+
15+
fn test_db() -> Arc<api::VectorDb> {
16+
let temp_dir = tempfile::tempdir().unwrap();
17+
let db = api::init_api(DbConfig {
18+
storage_type: StorageType::InMemory,
19+
index_type: IndexType::Flat,
20+
data_path: temp_dir.path().to_path_buf(),
21+
dimension: 3,
22+
similarity: Similarity::Cosine,
23+
hnsw_config: HnswConfig::default(),
24+
kd_tree_config: KDTreeConfig::default(),
25+
})
26+
.unwrap();
27+
Arc::new(db)
28+
}
29+
30+
fn test_server() -> TestServer {
31+
let keys = ApiKeyStore::new(vec![
32+
ApiKeyEntry {
33+
name: "full".to_string(),
34+
role: ApiKeyRole::ReadWrite,
35+
key: API_KEY.to_string(),
36+
},
37+
ApiKeyEntry {
38+
name: "readonly".to_string(),
39+
role: ApiKeyRole::ReadOnly,
40+
key: READONLY_API_KEY.to_string(),
41+
},
42+
]);
43+
let router = create_router(test_db(), Arc::new(keys));
44+
TestServer::new(router).unwrap()
45+
}
46+
47+
#[tokio::test]
48+
async fn in_memory_storage_http_smoke_test() {
49+
let server = test_server();
50+
51+
let insert_response = server
52+
.post("/points")
53+
.add_header("api-key", API_KEY)
54+
.json(&json!({
55+
"vector": [1.0, 0.0, 0.0],
56+
"payload": {
57+
"content_type": "Text",
58+
"content": "smoke-test"
59+
}
60+
}))
61+
.await;
62+
insert_response.assert_status(StatusCode::CREATED);
63+
let insert_body: serde_json::Value = insert_response.json();
64+
let point_id = insert_body["point_id"].as_str().unwrap();
65+
66+
let get_response = server
67+
.get(&format!("/points/{point_id}"))
68+
.add_header("api-key", API_KEY)
69+
.await;
70+
get_response.assert_status_ok();
71+
let point_body: serde_json::Value = get_response.json();
72+
assert_eq!(point_body["payload"]["content"], "smoke-test");
73+
assert_eq!(point_body["vector"], json!([1.0, 0.0, 0.0]));
74+
75+
let search_response = server
76+
.post("/points/search")
77+
.add_header("api-key", API_KEY)
78+
.json(&json!({
79+
"vector": [1.0, 0.0, 0.0],
80+
"similarity": "Cosine",
81+
"limit": 1
82+
}))
83+
.await;
84+
search_response.assert_status_ok();
85+
let search_body: serde_json::Value = search_response.json();
86+
assert_eq!(search_body["results"], json!([point_id]));
87+
88+
let delete_response = server
89+
.delete(&format!("/points/{point_id}"))
90+
.add_header("api-key", API_KEY)
91+
.await;
92+
delete_response.assert_status(StatusCode::NO_CONTENT);
93+
94+
let missing_response = server
95+
.get(&format!("/points/{point_id}"))
96+
.add_header("api-key", API_KEY)
97+
.await;
98+
missing_response.assert_status(StatusCode::NOT_FOUND);
99+
}
100+
101+
#[tokio::test]
102+
async fn public_routes_require_no_key() {
103+
let server = test_server();
104+
server.get("/").await.assert_status_ok();
105+
server.get("/health").await.assert_status_ok();
106+
}
107+
108+
#[tokio::test]
109+
async fn protected_routes_reject_missing_key() {
110+
let server = test_server();
111+
server
112+
.get("/points/00000000-0000-0000-0000-000000000000")
113+
.await
114+
.assert_status(StatusCode::UNAUTHORIZED);
115+
server
116+
.post("/points")
117+
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}}))
118+
.await
119+
.assert_status(StatusCode::UNAUTHORIZED);
120+
}
121+
122+
#[tokio::test]
123+
async fn protected_routes_reject_wrong_key() {
124+
let server = test_server();
125+
server
126+
.get("/points/00000000-0000-0000-0000-000000000000")
127+
.add_header("api-key", "not-the-right-key")
128+
.await
129+
.assert_status(StatusCode::UNAUTHORIZED);
130+
}
131+
132+
#[tokio::test]
133+
async fn readonly_key_can_read_but_not_write() {
134+
let server = test_server();
135+
136+
let insert_response = server
137+
.post("/points")
138+
.add_header("api-key", API_KEY)
139+
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "x"}}))
140+
.await;
141+
let point_id = insert_response.json::<serde_json::Value>()["point_id"]
142+
.as_str()
143+
.unwrap()
144+
.to_string();
145+
146+
server
147+
.get(&format!("/points/{point_id}"))
148+
.add_header("api-key", READONLY_API_KEY)
149+
.await
150+
.assert_status_ok();
151+
152+
server
153+
.post("/points/search")
154+
.add_header("api-key", READONLY_API_KEY)
155+
.json(&json!({"vector": [1.0, 0.0, 0.0], "similarity": "Cosine", "limit": 1}))
156+
.await
157+
.assert_status_ok();
158+
159+
server
160+
.delete(&format!("/points/{point_id}"))
161+
.add_header("api-key", READONLY_API_KEY)
162+
.await
163+
.assert_status(StatusCode::FORBIDDEN);
164+
165+
server
166+
.post("/points")
167+
.add_header("api-key", READONLY_API_KEY)
168+
.json(&json!({"vector": [1.0, 0.0, 0.0], "payload": {"content_type": "Text", "content": "y"}}))
169+
.await
170+
.assert_status(StatusCode::FORBIDDEN);
171+
}

‎crates/index/src/flat/constants.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
use defs::Magic;
2+
3+
pub const FLAT_MAGIC_BYTES: Magic = [0x00, 0x00, 0x00, 0x01];

‎crates/index/src/flat/mod.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
1-
use defs::Magic;
2-
1+
pub mod constants;
32
pub mod index;
43
mod serialize;
54

65
#[cfg(test)]
76
mod tests;
87

9-
pub const FLAT_MAGIC_BYTES: Magic = [0x00, 0x00, 0x00, 0x01];
8+
pub use constants::FLAT_MAGIC_BYTES;

‎crates/index/src/hnsw/constants.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
use defs::Magic;
2+
3+
pub const HNSW_MAGIC_BYTES: Magic = [0x02, 0x01, 0x03, 0x00];

0 commit comments

Comments
 (0)