tests for redisclient

This commit is contained in:
Per Stark
2024-10-02 11:17:42 +02:00
parent 8e7e762130
commit 7d13c06a51
8 changed files with 446 additions and 214 deletions

View File

@@ -13,10 +13,11 @@ use thiserror::Error;
use tracing::info;
use uuid::Uuid;
use crate::redis::client::RedisClient;
use crate::redis::client::{RedisClient, RedisClientTrait};
/// Represents metadata and storage information for a file.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]
pub struct FileInfo {
pub uuid: Uuid,
pub sha256: String,

View File

@@ -7,6 +7,7 @@ use crate::redis::client::RedisClient;
use super::{file_info::FileInfo, ingress_object::IngressObject };
/// Struct defining the expected body when ingressing content.
#[derive(Serialize, Deserialize, Debug)]
pub struct IngressInput {
@@ -107,3 +108,4 @@ pub async fn create_ingress_objects(
Ok(object_list)
}

View File

@@ -1,3 +1,4 @@
use axum::async_trait;
use redis::AsyncCommands;
use thiserror::Error;
use uuid::Uuid;
@@ -16,6 +17,32 @@ pub enum RedisError {
// Add more error variants as needed.
}
/// Defines the behavior for Redis client operations.
#[cfg_attr(test, mockall::automock)]
#[async_trait]
pub trait RedisClientTrait: Send + Sync {
/// Establishes a new multiplexed asynchronous connection to Redis.
async fn get_connection(&self) -> Result<redis::aio::MultiplexedConnection, RedisError>;
/// Stores `FileInfo` in Redis using SHA256 as the key.
async fn set_file_info(&self, sha256: &str, file_info: &FileInfo) -> Result<(), RedisError>;
/// Retrieves `FileInfo` from Redis using SHA256 as the key.
async fn get_file_info_by_sha(&self, sha256: &str) -> Result<Option<FileInfo>, RedisError>;
/// Deletes `FileInfo` from Redis using SHA256 as the key.
async fn delete_file_info(&self, sha256: &str) -> Result<(), RedisError>;
/// Sets a mapping from UUID to SHA256.
async fn set_sha_uuid_mapping(&self, uuid: &Uuid, sha256: &str) -> Result<(), RedisError>;
/// Retrieves the SHA256 hash associated with a given UUID.
async fn get_sha_by_uuid(&self, uuid: &Uuid) -> Result<Option<String>, RedisError>;
/// Deletes the UUID to SHA256 mapping from Redis.
async fn delete_sha_uuid_mapping(&self, uuid: &Uuid) -> Result<(), RedisError>;
}
/// Provides Redis-related operations for `FileInfo`.
pub struct RedisClient {
redis_url: String,
@@ -36,13 +63,15 @@ impl RedisClient {
redis_url: redis_url.to_string(),
}
}
}
#[async_trait]
impl RedisClientTrait for RedisClient {
/// Establishes a new multiplexed asynchronous connection to Redis.
///
/// # Returns
///
/// * `MultiplexedConnection` - The established connection.
pub async fn get_connection(&self) -> Result<redis::aio::MultiplexedConnection, RedisError> {
async fn get_connection(&self) -> Result<redis::aio::MultiplexedConnection, RedisError> {
let client = redis::Client::open(self.redis_url.clone())
.map_err(|e| RedisError::ConnectionError(e.to_string()))?;
let conn = client
@@ -55,14 +84,12 @@ impl RedisClient {
/// Stores `FileInfo` in Redis using SHA256 as the key.
///
/// # Arguments
///
/// * `sha256` - The SHA256 hash of the file.
/// * `file_info` - The `FileInfo` object to store.
///
/// # Returns
///
/// * `Result<(), RedisError>` - Empty result or an error.
pub async fn set_file_info(&self, sha256: &str, file_info: &FileInfo) -> Result<(), RedisError> {
async fn set_file_info(&self, sha256: &str, file_info: &FileInfo) -> Result<(), RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("file_info:{}", sha256);
let value = serde_json::to_string(file_info)
@@ -75,13 +102,11 @@ impl RedisClient {
/// Retrieves `FileInfo` from Redis using SHA256 as the key.
///
/// # Arguments
///
/// * `sha256` - The SHA256 hash of the file.
///
/// # Returns
///
/// * `Result<Option<FileInfo>, RedisError>` - The `FileInfo` if found, otherwise `None`, or an error.
pub async fn get_file_info_by_sha(&self, sha256: &str) -> Result<Option<FileInfo>, RedisError> {
async fn get_file_info_by_sha(&self, sha256: &str) -> Result<Option<FileInfo>, RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("file_info:{}", sha256);
let value: Option<String> = conn.get(key).await
@@ -98,13 +123,11 @@ impl RedisClient {
/// Deletes `FileInfo` from Redis using SHA256 as the key.
///
/// # Arguments
///
/// * `sha256` - The SHA256 hash of the file.
///
/// # Returns
///
/// * `Result<(), RedisError>` - Empty result or an error.
pub async fn delete_file_info(&self, sha256: &str) -> Result<(), RedisError> {
async fn delete_file_info(&self, sha256: &str) -> Result<(), RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("file_info:{}", sha256);
conn.del(key).await
@@ -115,14 +138,12 @@ impl RedisClient {
/// Sets a mapping from UUID to SHA256.
///
/// # Arguments
///
/// * `uuid` - The UUID of the file.
/// * `sha256` - The SHA256 hash of the file.
///
/// # Returns
///
/// * `Result<(), RedisError>` - Empty result or an error.
pub async fn set_sha_uuid_mapping(&self, uuid: &Uuid, sha256: &str) -> Result<(), RedisError> {
async fn set_sha_uuid_mapping(&self, uuid: &Uuid, sha256: &str) -> Result<(), RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("uuid_sha:{}", uuid);
conn.set(key, sha256).await
@@ -133,13 +154,11 @@ impl RedisClient {
/// Retrieves the SHA256 hash associated with a given UUID.
///
/// # Arguments
///
/// * `uuid` - The UUID of the file.
///
/// # Returns
///
/// * `Result<Option<String>, RedisError>` - The SHA256 hash if found, otherwise `None`, or an error.
pub async fn get_sha_by_uuid(&self, uuid: &Uuid) -> Result<Option<String>, RedisError> {
async fn get_sha_by_uuid(&self, uuid: &Uuid) -> Result<Option<String>, RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("uuid_sha:{}", uuid);
let sha: Option<String> = conn.get(key).await
@@ -156,7 +175,7 @@ impl RedisClient {
/// # Returns
///
/// * `Result<(), RedisError>` - Empty result or an error.
pub async fn delete_sha_uuid_mapping(&self, uuid: &Uuid) -> Result<(), RedisError> {
async fn delete_sha_uuid_mapping(&self, uuid: &Uuid) -> Result<(), RedisError> {
let mut conn = self.get_connection().await?;
let key = format!("uuid_sha:{}", uuid);
conn.del(key).await
@@ -164,3 +183,155 @@ impl RedisClient {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use mockall::predicate::*;
use uuid::Uuid;
#[tokio::test]
async fn test_set_file_info() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_sha = "dummysha256hash".to_string();
let test_file_info = FileInfo {
uuid: Uuid::new_v4(),
sha256: test_sha.clone(),
path: "/path/to/file".to_string(),
mime_type: "text/plain".to_string(),
};
// Setup expectation for `set_file_info`.
mock_redis
.expect_set_file_info()
.with(eq(test_sha.clone()), eq(test_file_info.clone()))
.times(1)
.returning(|_, _| Ok(()) );
// Call `set_file_info` on the mock.
let set_result = mock_redis.set_file_info(&test_sha, &test_file_info).await;
assert!(set_result.is_ok());
}
#[tokio::test]
async fn test_get_file_info_by_sha() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_sha = "dummysha256hash".to_string();
let test_file_info = FileInfo {
uuid: Uuid::new_v4(),
sha256: test_sha.clone(),
path: "/path/to/file".to_string(),
mime_type: "text/plain".to_string(),
};
// Clone the FileInfo to use inside the closure.
let fi_clone = test_file_info.clone();
// Setup expectation for `get_file_info_by_sha`.
mock_redis
.expect_get_file_info_by_sha()
.with(eq(test_sha.clone()))
.times(1)
.returning(move |_: &str| {
// Return the cloned FileInfo.
let fi_inner = fi_clone.clone();
Ok(Some(fi_inner))
});
// Call `get_file_info_by_sha` on the mock.
let get_result = mock_redis.get_file_info_by_sha(&test_sha).await;
assert!(get_result.is_ok());
assert_eq!(get_result.unwrap(), Some(test_file_info));
}
#[tokio::test]
async fn test_delete_file_info() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_sha = "dummysha256hash".to_string();
// Setup expectation for `delete_file_info`.
mock_redis
.expect_delete_file_info()
.with(eq(test_sha.clone()))
.times(1)
.returning(|_: &str| Ok(()) );
// Call `delete_file_info` on the mock.
let delete_result = mock_redis.delete_file_info(&test_sha).await;
assert!(delete_result.is_ok());
}
#[tokio::test]
async fn test_set_sha_uuid_mapping() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_uuid = Uuid::new_v4();
let test_sha = "dummysha256hash".to_string();
// Setup expectation for `set_sha_uuid_mapping`.
mock_redis
.expect_set_sha_uuid_mapping()
.with(eq(test_uuid.clone()), eq(test_sha.clone()))
.times(1)
.returning(|_, _| Ok(()) );
// Call `set_sha_uuid_mapping` on the mock.
let set_result = mock_redis.set_sha_uuid_mapping(&test_uuid, &test_sha).await;
assert!(set_result.is_ok());
}
#[tokio::test]
async fn test_get_sha_by_uuid() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_uuid = Uuid::new_v4();
let test_sha = "dummysha256hash".to_string();
// Clone the SHA to use inside the closure.
let sha_clone = test_sha.clone();
// Setup expectation for `get_sha_by_uuid`.
mock_redis
.expect_get_sha_by_uuid()
.with(eq(test_uuid.clone()))
.times(1)
.returning(move |_: &Uuid| {
let sha_inner = sha_clone.clone();
Ok(Some(sha_inner))
});
// Call `get_sha_by_uuid` on the mock.
let get_result = mock_redis.get_sha_by_uuid(&test_uuid).await;
assert!(get_result.is_ok());
assert_eq!(get_result.unwrap(), Some(test_sha));
}
#[tokio::test]
async fn test_delete_sha_uuid_mapping() {
// Initialize the mock.
let mut mock_redis = MockRedisClientTrait::new();
let test_uuid = Uuid::new_v4();
// Setup expectation for `delete_sha_uuid_mapping`.
mock_redis
.expect_delete_sha_uuid_mapping()
.with(eq(test_uuid.clone()))
.times(1)
.returning(|_: &Uuid| Ok(()) );
// Call `delete_sha_uuid_mapping` on the mock.
let delete_result = mock_redis.delete_sha_uuid_mapping(&test_uuid).await;
assert!(delete_result.is_ok());
}
}

View File

@@ -1,192 +0,0 @@
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::info;
use url::Url;
use uuid::Uuid;
use std::path::Path;
use tokio::fs;
/// Struct to reference stored files.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Reference {
pub uuid: Uuid,
pub path: String,
}
impl Reference {
/// Creates a new Reference with a generated UUID.
pub fn new(path: String) -> Self {
Self {
uuid: Uuid::new_v4(),
path,
}
}
}
/// Enum representing different types of content.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum Content {
Text(String),
Url(String),
Document(Reference),
Video(Reference),
Audio(Reference),
// Extend with more variants as needed
}
impl Content {
/// Retrieves the path from a reference if the content is a Reference variant.
pub fn get_path(&self) -> Option<&str> {
match self {
Content::Document(ref r) | Content::Video(ref r) | Content::Audio(ref r) => Some(&r.path),
_ => None,
}
}
}
#[derive(Error, Debug)]
pub enum IngressContentError {
#[error("IO error occurred: {0}")]
Io(#[from] std::io::Error),
#[error("UTF-8 conversion error: {0}")]
Utf8(#[from] std::string::FromUtf8Error),
#[error("MIME type detection failed for input: {0}")]
MimeDetection(String),
#[error("Unsupported MIME type: {0}")]
UnsupportedMime(String),
#[error("URL parse error: {0}")]
UrlParse(#[from] url::ParseError),
// Add more error variants as needed.
}
#[derive(Serialize, Deserialize, Debug)]
pub struct IngressContent {
pub content: Content,
pub category: String,
pub instructions: String,
}
impl IngressContent {
/// Creates a new IngressContent instance from the given input.
///
/// # Arguments
///
/// * `input` - A string slice that holds the input content, which can be text, a file path, or a URL.
/// * `category` - A string slice representing the category of the content.
/// * `instructions` - A string slice containing instructions for processing the content.
///
/// # Returns
///
/// * `Result<IngressContent, IngressContentError>` - The result containing either the IngressContent instance or an error.
pub async fn new(
input: &str,
category: &str,
instructions: &str,
) -> Result<IngressContent, IngressContentError> {
// Check if the input is a valid URL
if let Ok(url) = Url::parse(input) {
info!("Detected URL: {}", url);
return Ok(IngressContent {
content: Content::Url(url.to_string()),
category: category.to_string(),
instructions: instructions.to_string(),
});
}
// Attempt to treat the input as a file path
if let Ok(metadata) = tokio::fs::metadata(input).await {
if metadata.is_file() {
info!("Processing as file path: {}", input);
let mime = mime_guess::from_path(input).first_or(mime::TEXT_PLAIN);
let reference = Self::store_file(input, &mime).await?;
let content = match mime.type_() {
mime::TEXT | mime::APPLICATION => Content::Document(reference),
mime::VIDEO => Content::Video(reference),
mime::AUDIO => Content::Audio(reference),
other => {
info!("Detected unsupported MIME type: {}", other);
return Err(IngressContentError::UnsupportedMime(mime.to_string()));
}
};
return Ok(IngressContent {
content,
category: category.to_string(),
instructions: instructions.to_string(),
});
}
}
// Treat the input as plain text if it's neither a URL nor a file path
info!("Treating input as plain text");
Ok(IngressContent {
content: Content::Text(input.to_string()),
category: category.to_string(),
instructions: instructions.to_string(),
})
}
/// Stores the file into 'data/' directory and returns a Reference.
async fn store_file(input_path: &str, mime: &mime::Mime) -> Result<Reference, IngressContentError> {
return Ok(Reference::new(input_path.to_string()));
// Define the data directory
let data_dir = Path::new("data/");
// Ensure 'data/' directory exists; create it if it doesn't
fs::create_dir_all(data_dir).await.map_err(IngressContentError::Io)?;
// Generate a UUID for the file
let uuid = Uuid::new_v4();
// Determine the file extension based on MIME type
// let extension = Some(mime_guess::get_mime_extensions(mime)).unwrap_or("bin");
// Create a unique filename using UUID and extension
let file_name = format!("{}.{}", uuid, extension);
// Define the full file path
let file_path = data_dir.join(&file_name);
// Copy the original file to the 'data/' directory with the new filename
fs::copy(input_path, &file_path).await.map_err(IngressContentError::Io)?;
// Return a new Reference
Ok(Reference::new(file_path.to_string_lossy().to_string()))
}
/// Example method to handle content. Implement your actual logic here.
pub fn handle_content(&self) {
match &self.content {
Content::Text(text) => {
// Handle text content
println!("Text: {}", text);
}
Content::Url(url) => {
// Handle URL content
println!("URL: {}", url);
}
Content::Document(ref reference) => {
// Handle Document content via reference
println!("Document Reference: UUID: {}, Path: {}", reference.uuid, reference.path);
// Optionally, read the file from reference.path
}
Content::Video(ref reference) => {
// Handle Video content via reference
println!("Video Reference: UUID: {}, Path: {}", reference.uuid, reference.path);
// Optionally, read the file from reference.path
}
Content::Audio(ref reference) => {
// Handle Audio content via reference
println!("Audio Reference: UUID: {}, Path: {}", reference.uuid, reference.path);
// Optionally, read the file from reference.path
}
// Handle additional content types
}
}
}

View File

@@ -1,2 +1 @@
// pub mod mime;
// pub mod llm;