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 has_distinct_sampling_key(&self) -> bool {
21        self.own_key != self.sampling_key
22    }
23
24    pub fn from_envelope(envelope: &Envelope) -> Self {
25        let own_key = envelope.meta().public_key();
26        let sampling_key = envelope.sampling_key().unwrap_or(own_key);
27
28        Self::new(own_key, sampling_key)
29    }
30
31    pub fn iter(&self) -> impl Iterator<Item = ProjectKey> {
32        let Self {
33            own_key,
34            sampling_key,
35        } = self;
36
37        std::iter::once(*own_key).chain((own_key != sampling_key).then_some(*sampling_key))
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44    use std::collections::HashSet;
45
46    #[test]
47    fn test_project_key_pair_new() {
48        let own = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
49        let sampling = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
50
51        let pair = ProjectKeyPair::new(own, sampling);
52        assert_eq!(pair.own_key, own);
53        assert_eq!(pair.sampling_key, sampling);
54    }
55
56    #[test]
57    fn test_project_key_pair_equality() {
58        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
59        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
60
61        let pair1 = ProjectKeyPair::new(key1, key2);
62        let pair2 = ProjectKeyPair::new(key1, key2);
63        let pair3 = ProjectKeyPair::new(key2, key1);
64
65        assert_eq!(pair1, pair2);
66        assert_ne!(pair1, pair3);
67    }
68
69    #[test]
70    fn test_project_key_pair_ordering() {
71        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
72        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
73
74        let pair1 = ProjectKeyPair::new(key1, key2);
75        let pair2 = ProjectKeyPair::new(key2, key1);
76
77        assert!(pair1 < pair2);
78    }
79
80    #[test]
81    fn test_project_key_pair_hash() {
82        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
83        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
84
85        let pair1 = ProjectKeyPair::new(key1, key2);
86        let pair2 = ProjectKeyPair::new(key1, key2);
87        let pair3 = ProjectKeyPair::new(key2, key1);
88
89        let mut set = HashSet::new();
90        set.insert(pair1);
91        assert!(set.contains(&pair2));
92        assert!(!set.contains(&pair3));
93    }
94
95    #[test]
96    fn test_project_key_pair_iter() {
97        let key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
98        let key2 = ProjectKey::parse("b94ae32be2584e0bbd7a4cbb95971fee").unwrap();
99
100        // Test with different sampling key
101        let pair = ProjectKeyPair::new(key1, key2);
102        let keys: Vec<_> = pair.iter().collect();
103        assert_eq!(keys, vec![key1, key2]);
104
105        // Test with same key (should only yield one key)
106        let pair = ProjectKeyPair::new(key1, key1);
107        let keys: Vec<_> = pair.iter().collect();
108        assert_eq!(keys, vec![key1]);
109    }
110}
relay_server/services/buffer/common.rs

relay_server/services/buffer/
common.rs
