relay_server/services/buffer/
common.rs

1use relay_base_schema::project::ProjectKey;
2
3use crate::Envelope;
4
5/// Struct that represents two project keys.
6#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialOrd, PartialEq)]
7pub struct ProjectKeyPair {
8    pub own_key: ProjectKey,
9    pub sampling_key: ProjectKey,
10}
11
12impl ProjectKeyPair {
13    pub fn new(own_key: ProjectKey, sampling_key: ProjectKey) -> Self {
14        Self {
15            own_key,
16            sampling_key,
17        }
18    }
19
20    pub fn from_envelope(envelope: &Envelope) -> Self {
21        let own_key = envelope.meta().public_key();
22        let sampling_key = envelope.sampling_key().unwrap_or(own_key);
23
24        Self::new(own_key, sampling_key)
25    }
26
27    pub fn iter(&self) -> impl Iterator<Item = ProjectKey> {
28        let Self {
29            own_key,
30            sampling_key,
31        } = self;
32
33        std::iter::once(*own_key).chain((own_key != sampling_key).then_some(*sampling_key))
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use std::collections::HashSet;
41
42    #[test]
43    fn test_project_key_pair_new() {
44        let own = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
45        let sampling = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
46
47        let pair = ProjectKeyPair::new(own, sampling);
48        assert_eq!(pair.own_key, own);
49        assert_eq!(pair.sampling_key, sampling);
50    }
51
52    #[test]
53    fn test_project_key_pair_equality() {
54        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
55        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
56
57        let pair1 = ProjectKeyPair::new(key1, key2);
58        let pair2 = ProjectKeyPair::new(key1, key2);
59        let pair3 = ProjectKeyPair::new(key2, key1);
60
61        assert_eq!(pair1, pair2);
62        assert_ne!(pair1, pair3);
63    }
64
65    #[test]
66    fn test_project_key_pair_ordering() {
67        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
68        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
69
70        let pair1 = ProjectKeyPair::new(key1, key2);
71        let pair2 = ProjectKeyPair::new(key2, key1);
72
73        assert!(pair1 < pair2);
74    }
75
76    #[test]
77    fn test_project_key_pair_hash() {
78        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
79        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
80
81        let pair1 = ProjectKeyPair::new(key1, key2);
82        let pair2 = ProjectKeyPair::new(key1, key2);
83        let pair3 = ProjectKeyPair::new(key2, key1);
84
85        let mut set = HashSet::new();
86        set.insert(pair1);
87        assert!(set.contains(&pair2));
88        assert!(!set.contains(&pair3));
89    }
90
91    #[test]
92    fn test_project_key_pair_iter() {
93        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
94        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
95
96        // Test with different sampling key
97        let pair = ProjectKeyPair::new(key1, key2);
98        let keys: Vec<_> = pair.iter().collect();
99        assert_eq!(keys, vec![key1, key2]);
100
101        // Test with same key (should only yield one key)
102        let pair = ProjectKeyPair::new(key1, key1);
103        let keys: Vec<_> = pair.iter().collect();
104        assert_eq!(keys, vec![key1]);
105    }
106}
relay_server/services/buffer/common.rs

relay_server/services/buffer/
common.rs
