unleash-edge 0.2.0

Unleash edge is a proxy for Unleash. It can return both evaluated feature toggles as well as the raw data from Unleash's client API
Documentation
use crate::types::TokenRefresh;
use crate::types::{EdgeResult, EdgeToken};
use actix_web::http::header::EntityTag;
use async_trait::async_trait;
use dashmap::DashMap;
use unleash_types::client_features::ClientFeatures;

use super::repository::{DataSink, DataSource};

#[derive(Debug, Clone)]
pub struct MemoryProvider {
    data_store: DashMap<String, ClientFeatures>,
    token_store: DashMap<String, EdgeToken>,
    tokens_to_refresh: DashMap<String, TokenRefresh>,
}

fn key(token: &EdgeToken) -> String {
    token.environment.clone().unwrap()
}

impl Default for MemoryProvider {
    fn default() -> Self {
        Self::new()
    }
}

impl MemoryProvider {
    pub fn new() -> Self {
        Self {
            data_store: DashMap::new(),
            token_store: DashMap::new(),
            tokens_to_refresh: DashMap::new(),
        }
    }
}

#[async_trait]
impl DataSource for MemoryProvider {
    async fn get_tokens(&self) -> EdgeResult<Vec<EdgeToken>> {
        Ok(self.token_store.iter().map(|x| x.value().clone()).collect())
    }

    async fn get_token(&self, secret: &str) -> EdgeResult<Option<EdgeToken>> {
        Ok(self.token_store.get(secret).map(|x| x.clone()))
    }

    async fn get_refresh_tokens(&self) -> EdgeResult<Vec<TokenRefresh>> {
        Ok(self
            .tokens_to_refresh
            .iter()
            .map(|x| x.value().clone())
            .collect())
    }

    async fn get_client_features(&self, token: &EdgeToken) -> EdgeResult<Option<ClientFeatures>> {
        Ok(self.data_store.get(&key(token)).map(|v| v.value().clone()))
    }
}

#[async_trait]
impl DataSink for MemoryProvider {
    async fn sink_tokens(&self, tokens: Vec<EdgeToken>) -> EdgeResult<()> {
        for token in tokens {
            self.token_store.insert(token.token.clone(), token.clone());
        }
        Ok(())
    }

    async fn set_refresh_tokens(&self, tokens: Vec<&TokenRefresh>) -> EdgeResult<()> {
        self.tokens_to_refresh.clear();
        tokens.into_iter().for_each(|refresh| {
            self.tokens_to_refresh
                .insert(refresh.token.token.clone(), refresh.clone());
        });
        Ok(())
    }

    async fn sink_features(&self, token: &EdgeToken, features: ClientFeatures) -> EdgeResult<()> {
        self.data_store.insert(key(token), features);
        Ok(())
    }

    async fn update_last_check(&self, token: &EdgeToken) -> EdgeResult<()> {
        if let Some(mut token) = self.tokens_to_refresh.get_mut(&token.token) {
            token.last_check = Some(chrono::Utc::now());
        }
        Ok(())
    }

    async fn update_last_refresh(
        &self,
        token: &EdgeToken,
        etag: Option<EntityTag>,
    ) -> EdgeResult<()> {
        if let Some(mut token) = self.tokens_to_refresh.get_mut(&token.token) {
            token.last_check = Some(chrono::Utc::now());
            token.last_refreshed = Some(chrono::Utc::now());
            token.etag = etag;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::types::TokenValidationStatus;

    use super::*;

    #[tokio::test]
    async fn memory_provider_correctly_deduplicates_tokens() {
        let provider = MemoryProvider::default();
        provider
            .sink_tokens(vec![EdgeToken {
                token: "*:development.1d38eefdd7bf72676122b008dcf330f2f2aa2f3031438e1b7e8f0d1f"
                    .into(),
                ..EdgeToken::default()
            }])
            .await
            .unwrap();

        provider
            .sink_tokens(vec![EdgeToken {
                token: "*:development.1d38eefdd7bf72676122b008dcf330f2f2aa2f3031438e1b7e8f0d1f"
                    .into(),
                ..EdgeToken::default()
            }])
            .await
            .unwrap();

        assert!(provider.get_tokens().await.unwrap().len() == 1);
    }

    #[tokio::test]
    async fn memory_provider_correctly_determines_token_to_be_valid() {
        let provider = MemoryProvider::default();
        provider
            .sink_tokens(vec![EdgeToken {
                token: "*:development.1d38eefdd7bf72676122b008dcf330f2f2aa2f3031438e1b7e8f0d1f"
                    .into(),
                status: TokenValidationStatus::Validated,
                ..EdgeToken::default()
            }])
            .await
            .unwrap();

        assert_eq!(
            provider
                .get_token("*:development.1d38eefdd7bf72676122b008dcf330f2f2aa2f3031438e1b7e8f0d1f")
                .await
                .expect("Could not retrieve token details")
                .unwrap()
                .status,
            TokenValidationStatus::Validated
        )
    }
}