Unverified Commit 58da7cfe authored by Graham King's avatar Graham King Committed by GitHub
Browse files

test(key_value_store): Assert entries() works correctly (#4312)


Signed-off-by: default avatarGraham King <grahamk@nvidia.com>
parent 3cf67ee0
......@@ -471,3 +471,33 @@ fn a_to_fs_err(err: anyhow::Error) -> StoreError {
fn to_fs_err<E: std::error::Error>(err: E) -> StoreError {
StoreError::FilesystemError(err.to_string())
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use crate::storage::key_value_store::{
FileStore, Key, KeyValueBucket as _, KeyValueStore as _,
};
#[tokio::test]
async fn test_entries_full_path() {
let t = tempfile::tempdir().unwrap();
let m = FileStore::new(t.path());
let bucket = m.get_or_create_bucket("v1/tests", None).await.unwrap();
let _ = bucket
.insert(&Key::new("key1/multi/part"), "value1".into(), 0)
.await
.unwrap();
let _ = bucket
.insert(&Key::new("key2"), "value2".into(), 0)
.await
.unwrap();
let entries = bucket.entries().await.unwrap();
let keys: HashSet<String> = entries.into_keys().collect();
assert!(keys.contains("v1/tests/key1/multi/part"));
assert!(keys.contains("v1/tests/key2"));
}
}
......@@ -220,9 +220,36 @@ impl KeyValueBucket for MemoryBucketRef {
Some(bucket) => Ok(bucket
.data
.iter()
.map(|(k, (_rev, v))| (k.to_string(), v.clone()))
.map(|(k, (_rev, v))| ([self.name.clone(), k.to_string()].join("/"), v.clone()))
.collect()),
None => Err(StoreError::MissingBucket(self.name.clone())),
}
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use crate::storage::key_value_store::{
Key, KeyValueBucket as _, KeyValueStore as _, MemoryStore,
};
#[tokio::test]
async fn test_entries_full_path() {
let m = MemoryStore::new();
let bucket = m.get_or_create_bucket("bucket1", None).await.unwrap();
let _ = bucket
.insert(&Key::new("key1"), "value1".into(), 0)
.await
.unwrap();
let _ = bucket
.insert(&Key::new("key2"), "value2".into(), 0)
.await
.unwrap();
let entries = bucket.entries().await.unwrap();
let keys: HashSet<String> = entries.into_keys().collect();
assert!(keys.contains("bucket1/key1"));
assert!(keys.contains("bucket1/key2"));
}
}
......@@ -6,6 +6,7 @@
import logging
import os
import shutil
import tempfile
import numpy as np
import pytest
......@@ -90,6 +91,25 @@ def extract_params(param_map) -> dict:
return result
@pytest.fixture
def file_storage_backend():
"""Fixture that sets up and tears down file storage backend.
Creates a temporary directory for file-based KV storage and sets
the DYN_FILE_KV environment variable. Cleans up after the test.
"""
with tempfile.TemporaryDirectory() as tmpdir:
old_env = os.environ.get("DYN_FILE_KV")
os.environ["DYN_FILE_KV"] = tmpdir
logger.info(f"Set up file storage backend in: {tmpdir}")
yield tmpdir
# Cleanup
if old_env is not None:
os.environ["DYN_FILE_KV"] = old_env
else:
os.environ.pop("DYN_FILE_KV", None)
@pytest.mark.e2e
@pytest.mark.pre_merge
@pytest.mark.parametrize(
......@@ -101,7 +121,7 @@ def extract_params(param_map) -> dict:
],
ids=["no_params", "numeric_param", "mixed_params"],
)
def test_request_parameters(start_services, request_params):
def test_request_parameters(file_storage_backend, start_services, request_params):
"""Test gRPC request-level parameters are echoed through tensor models.
The worker acts as an identity function: echoes input tensors unchanged and
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment