rustrails_storage/service/

memory.rs

//! In-memory storage service used by tests and ephemeral flows.

use std::{sync::Arc, time::Duration};

use async_trait::async_trait;
use bytes::Bytes;
use dashmap::DashMap;
use url::Url;

use super::{StorageError, StorageService, checked_key};

/// Thread-safe in-memory storage backend.
#[derive(Debug, Clone)]
pub struct MemoryService {
    name: String,
    objects: Arc<DashMap<String, Bytes>>,
    base_url: Url,
}

impl MemoryService {
    /// Creates an empty in-memory service.
    ///
    /// # Errors
    ///
    /// Returns an error when the default base URL cannot be constructed.
    pub fn new(name: impl Into<String>) -> Result<Self, StorageError> {
        let base_url = Url::parse("memory://storage/")
            .map_err(|error| StorageError::InvalidUrl(error.to_string()))?;
        Ok(Self {
            name: name.into(),
            objects: Arc::new(DashMap::new()),
            base_url,
        })
    }

    /// Returns a copy with a custom public base URL.
    #[must_use]
    pub fn with_base_url(mut self, base_url: Url) -> Self {
        self.base_url = base_url;
        self
    }

    /// Returns the number of stored keys.
    #[must_use]
    pub fn len(&self) -> usize {
        self.objects.len()
    }

    /// Returns whether the service currently stores no keys.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.objects.is_empty()
    }
}

#[async_trait]
impl StorageService for MemoryService {
    fn name(&self) -> &str {
        &self.name
    }

    async fn upload(&self, key: &str, data: Bytes) -> Result<(), StorageError> {
        let key = checked_key(key)?.to_owned();
        if self.objects.contains_key(&key) {
            return Err(StorageError::DuplicateKey(key));
        }
        self.objects.insert(key, data);
        Ok(())
    }

    async fn download(&self, key: &str) -> Result<Bytes, StorageError> {
        let key = checked_key(key)?;
        self.objects
            .get(key)
            .map(|entry| entry.value().clone())
            .ok_or_else(|| StorageError::NotFound(key.to_owned()))
    }

    async fn delete(&self, key: &str) -> Result<(), StorageError> {
        let key = checked_key(key)?;
        let _ = self.objects.remove(key);
        Ok(())
    }

    async fn exists(&self, key: &str) -> Result<bool, StorageError> {
        let key = checked_key(key)?;
        Ok(self.objects.contains_key(key))
    }

    async fn url(&self, key: &str, expires_in: Duration) -> Result<Url, StorageError> {
        let key = checked_key(key)?;
        let mut url = self.base_url.clone();
        url.set_path(key);
        url.query_pairs_mut()
            .append_pair("service", &self.name)
            .append_pair("expires_in", &expires_in.as_secs().to_string());
        Ok(url)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn service() -> MemoryService {
        MemoryService::new("memory").expect("service should build")
    }

    #[tokio::test]
    async fn test_upload_and_download_round_trip() {
        let service = service();
        service
            .upload("a.txt", Bytes::from_static(b"hello"))
            .await
            .expect("upload should succeed");
        assert_eq!(
            service
                .download("a.txt")
                .await
                .expect("download should succeed"),
            Bytes::from_static(b"hello")
        );
    }

    #[tokio::test]
    async fn test_upload_rejects_duplicate_key() {
        let service = service();
        service
            .upload("a.txt", Bytes::from_static(b"one"))
            .await
            .expect("upload should succeed");
        let error = service
            .upload("a.txt", Bytes::from_static(b"two"))
            .await
            .expect_err("duplicate should fail");
        assert!(matches!(error, StorageError::DuplicateKey(key) if key == "a.txt"));
    }

    #[tokio::test]
    async fn test_download_missing_key_returns_not_found() {
        let service = service();
        let error = service
            .download("missing")
            .await
            .expect_err("download should fail");
        assert!(matches!(error, StorageError::NotFound(key) if key == "missing"));
    }

    #[tokio::test]
    async fn test_delete_removes_value() {
        let service = service();
        service
            .upload("a.txt", Bytes::from_static(b"hello"))
            .await
            .expect("upload should succeed");
        service
            .delete("a.txt")
            .await
            .expect("delete should succeed");
        assert!(
            !service
                .exists("a.txt")
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn test_delete_missing_key_is_a_noop() {
        let service = service();
        service
            .delete("missing")
            .await
            .expect("delete should succeed");
        assert!(service.is_empty());
    }

    #[tokio::test]
    async fn test_exists_tracks_presence() {
        let service = service();
        assert!(
            !service
                .exists("a.txt")
                .await
                .expect("exists should succeed")
        );
        service
            .upload("a.txt", Bytes::from_static(b"hello"))
            .await
            .expect("upload should succeed");
        assert!(
            service
                .exists("a.txt")
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn test_url_includes_expiry_and_service_name() {
        let service = service();
        let url = service
            .url("a.txt", Duration::from_secs(30))
            .await
            .expect("url should build");
        assert_eq!(
            url.as_str(),
            "memory://storage/a.txt?service=memory&expires_in=30"
        );
    }

    #[tokio::test]
    async fn test_zero_byte_upload_is_supported() {
        let service = service();
        service
            .upload("empty", Bytes::new())
            .await
            .expect("upload should succeed");
        assert_eq!(
            service
                .download("empty")
                .await
                .expect("download should succeed"),
            Bytes::new()
        );
    }

    #[tokio::test]
    async fn test_len_tracks_number_of_entries() {
        let service = service();
        service
            .upload("a", Bytes::from_static(b"a"))
            .await
            .expect("upload should succeed");
        service
            .upload("b", Bytes::from_static(b"b"))
            .await
            .expect("upload should succeed");
        assert_eq!(service.len(), 2);
    }

    #[tokio::test]
    async fn test_custom_base_url_is_respected() {
        let service = MemoryService::new("memory")
            .expect("service should build")
            .with_base_url(Url::parse("https://memory.example/").expect("url should parse"));
        let url = service
            .url("a.txt", Duration::from_secs(5))
            .await
            .expect("url should build");
        assert_eq!(
            url.as_str(),
            "https://memory.example/a.txt?service=memory&expires_in=5"
        );
    }

    #[tokio::test]
    async fn test_invalid_empty_key_rejected() {
        let service = service();
        let error = service
            .upload("", Bytes::from_static(b"hello"))
            .await
            .expect_err("upload should fail");
        assert!(matches!(error, StorageError::InvalidUrl(_)));
    }
}