this/events/sinks/
device_tokens.rs1use chrono::{DateTime, Utc};
13use serde::{Deserialize, Serialize};
14use std::collections::HashMap;
15use tokio::sync::RwLock;
16
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum Platform {
21 Ios,
23 Android,
25 Web,
27}
28
29impl std::fmt::Display for Platform {
30 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31 match self {
32 Platform::Ios => write!(f, "ios"),
33 Platform::Android => write!(f, "android"),
34 Platform::Web => write!(f, "web"),
35 }
36 }
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct DeviceToken {
42 pub token: String,
44
45 pub platform: Platform,
47
48 pub registered_at: DateTime<Utc>,
50}
51
52#[derive(Debug)]
57pub struct DeviceTokenStore {
58 tokens: RwLock<HashMap<String, Vec<DeviceToken>>>,
59}
60
61impl DeviceTokenStore {
62 pub fn new() -> Self {
64 Self {
65 tokens: RwLock::new(HashMap::new()),
66 }
67 }
68
69 pub async fn register(&self, user_id: &str, token: String, platform: Platform) {
74 let mut store = self.tokens.write().await;
75 let user_tokens = store.entry(user_id.to_string()).or_default();
76
77 if let Some(existing) = user_tokens.iter_mut().find(|t| t.token == token) {
79 existing.platform = platform;
80 existing.registered_at = Utc::now();
81 } else {
82 user_tokens.push(DeviceToken {
83 token,
84 platform,
85 registered_at: Utc::now(),
86 });
87 }
88 }
89
90 pub async fn unregister(&self, user_id: &str, token: &str) -> bool {
94 let mut store = self.tokens.write().await;
95 if let Some(user_tokens) = store.get_mut(user_id) {
96 let len_before = user_tokens.len();
97 user_tokens.retain(|t| t.token != token);
98 return user_tokens.len() < len_before;
99 }
100 false
101 }
102
103 pub async fn get_tokens(&self, user_id: &str) -> Vec<DeviceToken> {
105 let store = self.tokens.read().await;
106 store.get(user_id).cloned().unwrap_or_default()
107 }
108
109 pub async fn token_count(&self, user_id: &str) -> usize {
111 let store = self.tokens.read().await;
112 store.get(user_id).map(|t| t.len()).unwrap_or(0)
113 }
114
115 pub async fn remove_all(&self, user_id: &str) -> usize {
117 let mut store = self.tokens.write().await;
118 store.remove(user_id).map(|t| t.len()).unwrap_or(0)
119 }
120}
121
122impl Default for DeviceTokenStore {
123 fn default() -> Self {
124 Self::new()
125 }
126}
127
128#[cfg(test)]
129mod tests {
130 use super::*;
131
132 #[tokio::test]
133 async fn test_register_and_get_tokens() {
134 let store = DeviceTokenStore::new();
135
136 store
137 .register("user-A", "token-1".to_string(), Platform::Ios)
138 .await;
139 store
140 .register("user-A", "token-2".to_string(), Platform::Android)
141 .await;
142
143 let tokens = store.get_tokens("user-A").await;
144 assert_eq!(tokens.len(), 2);
145 assert_eq!(tokens[0].token, "token-1");
146 assert_eq!(tokens[0].platform, Platform::Ios);
147 assert_eq!(tokens[1].token, "token-2");
148 assert_eq!(tokens[1].platform, Platform::Android);
149 }
150
151 #[tokio::test]
152 async fn test_unregister() {
153 let store = DeviceTokenStore::new();
154
155 store
156 .register("user-A", "token-1".to_string(), Platform::Ios)
157 .await;
158 store
159 .register("user-A", "token-2".to_string(), Platform::Android)
160 .await;
161
162 assert_eq!(store.token_count("user-A").await, 2);
163
164 let removed = store.unregister("user-A", "token-1").await;
165 assert!(removed);
166 assert_eq!(store.token_count("user-A").await, 1);
167
168 let tokens = store.get_tokens("user-A").await;
169 assert_eq!(tokens[0].token, "token-2");
170 }
171
172 #[tokio::test]
173 async fn test_unregister_nonexistent() {
174 let store = DeviceTokenStore::new();
175 assert!(!store.unregister("user-A", "nonexistent").await);
176 }
177
178 #[tokio::test]
179 async fn test_register_duplicate_updates() {
180 let store = DeviceTokenStore::new();
181
182 store
183 .register("user-A", "token-1".to_string(), Platform::Ios)
184 .await;
185 store
187 .register("user-A", "token-1".to_string(), Platform::Android)
188 .await;
189
190 let tokens = store.get_tokens("user-A").await;
191 assert_eq!(tokens.len(), 1); assert_eq!(tokens[0].platform, Platform::Android); }
194
195 #[tokio::test]
196 async fn test_get_tokens_empty() {
197 let store = DeviceTokenStore::new();
198 assert!(store.get_tokens("nonexistent").await.is_empty());
199 assert_eq!(store.token_count("nonexistent").await, 0);
200 }
201
202 #[tokio::test]
203 async fn test_remove_all() {
204 let store = DeviceTokenStore::new();
205
206 store
207 .register("user-A", "token-1".to_string(), Platform::Ios)
208 .await;
209 store
210 .register("user-A", "token-2".to_string(), Platform::Android)
211 .await;
212
213 let removed = store.remove_all("user-A").await;
214 assert_eq!(removed, 2);
215 assert!(store.get_tokens("user-A").await.is_empty());
216 }
217
218 #[tokio::test]
219 async fn test_separate_users() {
220 let store = DeviceTokenStore::new();
221
222 store
223 .register("user-A", "token-a".to_string(), Platform::Ios)
224 .await;
225 store
226 .register("user-B", "token-b".to_string(), Platform::Android)
227 .await;
228
229 assert_eq!(store.token_count("user-A").await, 1);
230 assert_eq!(store.token_count("user-B").await, 1);
231
232 store.unregister("user-A", "token-a").await;
234 assert_eq!(store.token_count("user-A").await, 0);
235 assert_eq!(store.token_count("user-B").await, 1);
236 }
237}