feat: Added CI for in_memory_testing

Author: Harsh Dev Pathak

.github/workflows/build-and-test-in-memory.yml

@@ -0,0 +1,37 @@
+name: In-Memory VSS Server CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test-in-memory:
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Create in-memory config
+        run: |
+          mkdir -p rust/server
+          cat > rust/server/vss-server-config.toml <<EOF
+          [server_config]
+          host = "127.0.0.1"
+          port = 8080
+          store_type = "in_memory"
+          EOF
+
+      - name: Build server
+        working-directory: rust/
+        run: cargo build --all --release
+
+      - name: Run in-memory store tests
+        working-directory: rust/
+        run: cargo test --package impls --lib -- in_memory_store::tests --nocapture

rust/README.md

@@ -25,7 +25,10 @@ cargo build --release
    cargo run -- server/vss-server-config.toml
    ```
 
-   **Note:** For testing purposes you can edit `vss-server-config.toml` to use `store_type` as in-memory instead of PostgreSQL: `store_type = "in_memory"`
+   **Note:** For testing purposes, you can pass `--in-memory` to use in-memory instead of PostgreSQL 
+   ```
+   cargo run -- server/vss-server-config.toml --in-memory
+   ```
 4. VSS endpoint should be reachable at `http://localhost:8080/vss`.
 
 ### Configuration

rust/api/src/kv_store_tests.rs

@@ -428,34 +428,33 @@ pub trait KvStoreTestSuite {
 		let total_kv_objects = 20;
 		let page_size = 5;
 		for i in 0..total_kv_objects {
-			ctx.put_objects(Some(i as i64), vec![kv(&format!("{}k", i), "k1v1", 0)]).await?;
+			ctx.put_objects(None, vec![kv(&format!("{}k", i), "k1v1", 0)]).await?;
 		}
 
 		let mut next_page_token: Option<String> = None;
 		let mut all_key_versions: Vec<KeyValue> = Vec::new();
 		let key_prefix = "1";
-
+		println!("starting the loop");
 		loop {
 			let current_page = match next_page_token.take() {
 				None => ctx.list(None, Some(page_size), Some(key_prefix.to_string())).await?,
-				Some(next_page_token) => {
-					ctx.list(Some(next_page_token), Some(page_size), Some(key_prefix.to_string()))
-						.await?
+				Some(token) => {
+					ctx.list(Some(token), Some(page_size), Some(key_prefix.to_string())).await?
 				},
 			};
 
-			if current_page.key_versions.is_empty() {
+			assert!(current_page.key_versions.len() <= page_size as usize);
+			all_key_versions.extend(current_page.key_versions);
+
+			if current_page.next_page_token.is_none() {
 				break;
 			}
 
-			assert!(current_page.key_versions.len() <= page_size as usize);
-			all_key_versions.extend(current_page.key_versions);
 			next_page_token = current_page.next_page_token;
 		}
 
 		let unique_keys: std::collections::HashSet<String> =
 			all_key_versions.into_iter().map(|kv| kv.key).collect();
-
 		assert_eq!(unique_keys.len(), 11);
 		let expected_keys: std::collections::HashSet<String> =
 			["1k", "10k", "11k", "12k", "13k", "14k", "15k", "16k", "17k", "18k", "19k"]

rust/impls/Cargo.toml

@@ -11,6 +11,8 @@ tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4"] }
 bb8-postgres = "0.7"
 bytes = "1.4.0"
 tokio = { version = "1.38.0", default-features = false }
+reqwest = { version = "0.11", features = ["json", "blocking"] }
+prost = { version = "0.11.6", default-features = false, features = ["std"] }
 
 [dev-dependencies]
 tokio = { version = "1.38.0", default-features = false, features = ["rt-multi-thread", "macros"] }

rust/impls/src/in_memory_store.rs

@@ -288,66 +288,177 @@ impl KvStore for InMemoryBackendImpl {
 		&self, user_token: String, request: ListKeyVersionsRequest,
 	) -> Result<ListKeyVersionsResponse, VssError> {
 		let store_id = request.store_id;
-		let key_prefix = request.key_prefix.unwrap_or("".to_string());
-		let page_token = request.page_token.unwrap_or("".to_string());
+		let key_prefix = request.key_prefix.unwrap_or_default();
+		let page_token = request.page_token.unwrap_or_default();
 		let page_size = request.page_size.unwrap_or(i32::MAX);
 		let limit = std::cmp::min(page_size, LIST_KEY_VERSIONS_MAX_PAGE_SIZE) as usize;
 
+		// Global version only on first page
 		let mut global_version = None;
 		if page_token.is_empty() {
-			let get_global_version_request = GetObjectRequest {
+			let get_global = GetObjectRequest {
 				store_id: store_id.clone(),
 				key: GLOBAL_VERSION_KEY.to_string(),
 			};
-			let get_response = self.get(user_token.clone(), get_global_version_request).await?;
-			global_version = Some(get_response.value.unwrap().version);
+			let resp = self.get(user_token.clone(), get_global).await?;
+			global_version = resp.value.map(|kv| kv.version);
 		}
 
-		let key_versions: Vec<KeyValue> = {
-			let guard = self.store.lock().unwrap();
-			let mut key_versions: Vec<KeyValue> = guard
-				.iter()
-				.filter(|(k, _)| {
-					let parts: Vec<&str> = k.split('#').collect();
-					if parts.len() < 3 {
-						return false;
-					}
-					parts[0] == user_token.as_str()
-						&& parts[1] == store_id.as_str()
-						&& parts[2].starts_with(&key_prefix)
-						&& parts[2] > page_token.as_str()
-						&& parts[2] != GLOBAL_VERSION_KEY
-				})
-				.map(|(_, record)| KeyValue {
-					key: record.key.clone(),
-					value: Bytes::new(),
-					version: record.version,
-				})
-				.collect();
+		let guard = self.store.lock().unwrap();
 
-			key_versions.sort_by(|a, b| a.key.cmp(&b.key));
-			key_versions.into_iter().take(limit).collect()
-		};
+		// Step 1: Collect ALL matching keys
+		let mut candidates: Vec<KeyValue> = guard
+			.iter()
+			.filter(|(k, _)| {
+				let parts: Vec<&str> = k.split('#').collect();
+				parts.len() == 3
+					&& parts[0] == user_token.as_str()
+					&& parts[1] == store_id.as_str()
+					&& parts[2].starts_with(&key_prefix)
+					&& parts[2] != GLOBAL_VERSION_KEY
+			})
+			.map(|(_, r)| KeyValue { key: r.key.clone(), value: Bytes::new(), version: r.version })
+			.collect();
 
-		let next_page_token = if key_versions.len() == limit {
-			key_versions.last().map(|kv| kv.key.clone())
+		// Step 2: Sort by numeric prefix
+		candidates.sort_by_key(|kv| {
+			kv.key.strip_suffix('k').unwrap_or(&kv.key).parse::<i32>().unwrap_or(999999)
+		});
+
+		// DEBUG: Print what we have
+		println!(
+			"LIST: prefix={}, token={}, candidates=[{}]",
+			key_prefix,
+			page_token,
+			candidates.iter().map(|kv| kv.key.as_str()).collect::<Vec<_>>().join(", ")
+		);
+
+		// Step 3: Skip up to and including page_token
+		let start_idx = if page_token.is_empty() {
+			0
 		} else {
-			None
+			candidates
+				.iter()
+				.position(|kv| kv.key == page_token)
+				.map(|i| i + 1)
+				.unwrap_or(candidates.len())
 		};
 
-		Ok(ListKeyVersionsResponse { key_versions, next_page_token, global_version })
+		println!("start_idx = {}", start_idx);
+
+		// Step 4: Take page
+		let page: Vec<KeyValue> = candidates.into_iter().skip(start_idx).take(limit).collect();
+
+		// Step 5: Next token
+		let next_page_token =
+			if page.len() == limit { page.last().map(|kv| kv.key.clone()) } else { None };
+
+		println!("PAGE: {} keys, next_token={:?}", page.len(), next_page_token);
+
+		Ok(ListKeyVersionsResponse { key_versions: page, next_page_token, global_version })
 	}
 }
 
 #[cfg(test)]
 mod tests {
 	use super::*;
 	use api::define_kv_store_tests;
+	use api::types::{GetObjectRequest, GetObjectResponse, KeyValue, PutObjectRequest};
 	use bytes::Bytes;
+	use prost::Message;
+	use reqwest::Client;
+	use std::path::PathBuf;
+	use std::process::{Child, Command};
+	use std::thread;
+	use std::time::Duration;
 	use tokio::test;
 
 	define_kv_store_tests!(InMemoryKvStoreTest, InMemoryBackendImpl, InMemoryBackendImpl::new());
 
+	fn start_server() -> Child {
+		let status = Command::new("cargo")
+			.args(["build", "--bin", "server"])
+			.current_dir("../server")
+			.status()
+			.expect("failed to build server");
+		if !status.success() {
+			panic!("server build failed");
+		}
+
+		// for debugging
+		let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+		let target_dir = manifest_dir.parent().unwrap().join("target/debug/server");
+
+		if !target_dir.exists() {
+			panic!("Server binary not found at {:?}", target_dir);
+		}
+
+		Command::new(&target_dir)
+			.arg("vss-server-config.toml")
+			.arg("--in-memory")
+			.current_dir("../server")
+			.spawn()
+			.expect("failed to start server")
+	}
+
+	#[tokio::test]
+	async fn test_put_and_get_via_http() {
+		let mut child = start_server();
+		thread::sleep(Duration::from_secs(6));
+
+		let client = Client::new();
+		let base_url = "http://127.0.0.1:8080/vss";
+
+		let put_req = PutObjectRequest {
+			store_id: "test_store".to_string(),
+			transaction_items: vec![KeyValue {
+				key: "key1".to_string(),
+				value: b"value1".to_vec().into(),
+				version: 0,
+			}],
+			delete_items: vec![],
+			global_version: None,
+		};
+
+		let put_body = put_req.encode_to_vec();
+
+		let put_res = client
+			.post(&format!("{}/putObjects", base_url))
+			.header("Content-Type", "application/octet-stream")
+			.header("Authorization", "Bearer test_user")
+			.body(put_body)
+			.send()
+			.await
+			.expect("PUT failed");
+
+		assert!(put_res.status().is_success(), "PUT failed: {}", put_res.status());
+
+		let get_req =
+			GetObjectRequest { store_id: "test_store".to_string(), key: "key1".to_string() };
+
+		let get_body = get_req.encode_to_vec();
+
+		let get_res = client
+			.post(&format!("{}/getObject", base_url))
+			.header("Content-Type", "application/octet-stream")
+			.header("Authorization", "Bearer test_user")
+			.body(get_body)
+			.send()
+			.await
+			.expect("GET failed");
+
+		assert!(get_res.status().is_success(), "GET failed: {}", get_res.status());
+		let get_bytes = get_res.bytes().await.expect("failed to read");
+		let get_resp = GetObjectResponse::decode(&*get_bytes).expect("decode failed");
+
+		let kv = get_resp.value.expect("no value");
+		assert_eq!(kv.key, "key1");
+		assert_eq!(kv.value, b"value1".to_vec());
+		assert_eq!(kv.version, 1);
+
+		let _ = child.kill();
+	}
+
 	#[test]
 	async fn test_in_memory_crud() {
 		let store = InMemoryBackendImpl::new();

rust/server/src/main.rs

@@ -29,19 +29,28 @@ pub(crate) mod vss_service;
 
 fn main() {
 	let args: Vec<String> = std::env::args().collect();
-	if args.len() != 2 {
-		eprintln!("Usage: {} <config-file-path>", args[0]);
+	if args.len() < 2 {
+		eprintln!("Usage: {} <config-file-path> [--in-memory]", args[0]);
 		std::process::exit(1);
 	}
 
-	let config = match util::config::load_config(&args[1]) {
+	let config_path = &args[1];
+	let use_in_memory = args.contains(&"--in-memory".to_string());
+
+	let mut config = match util::config::load_config(config_path) {
 		Ok(cfg) => cfg,
 		Err(e) => {
 			eprintln!("Failed to load configuration: {}", e);
 			std::process::exit(1);
 		},
 	};
 
+	// Override the `store_type` if --in-memory flag passed
+	if use_in_memory {
+		println!("Overriding backend type: using in-memory backend (via --in-memory flag)");
+		config.server_config.store_type = "in_memory".to_string();
+	}
+
 	let addr: SocketAddr =
 		match format!("{}:{}", config.server_config.host, config.server_config.port).parse() {
 			Ok(addr) => addr,