Implement end-to-end tests for TC token

src/features/presence.rs

@@ -140,6 +140,27 @@ impl<'a> Presence<'a> {
         Ok(())
     }
 
+    /// Re-subscribe presence if the JID has an active subscription.
+    /// Does not modify the tracking set.
+    pub(crate) async fn re_subscribe_when_active(&self, jid: &Jid) -> Result<(), anyhow::Error> {
+        if !self
+            .client
+            .presence_subscriptions
+            .lock()
+            .await
+            .contains(jid)
+        {
+            return Ok(());
+        }
+
+        let node = self.build_subscription_node(jid).await;
+        self.client
+            .send_node(node)
+            .await
+            .map_err(anyhow::Error::from)?;
+        Ok(())
+    }
+
     /// Unsubscribe from a contact's presence updates.
     ///
     /// Sends a `<presence type="unsubscribe">` stanza to the target JID.
@@ -204,14 +225,7 @@ impl Client {
                 return;
             }
 
-            // Check membership before re-subscribing — a concurrent unsubscribe()
-            // call may have removed this JID while we were iterating.
-            if !self.presence_subscriptions.lock().await.contains(&jid) {
-                debug!("Skipping re-subscribe for {jid}: unsubscribed during iteration");
-                continue;
-            }
-
-            if let Err(err) = self.presence().subscribe(&jid).await {
+            if let Err(err) = self.presence().re_subscribe_when_active(&jid).await {
                 warn!("Failed to re-subscribe to presence for {jid}: {err:?}");
             }
         }

src/handlers/notification.rs

@@ -546,7 +546,9 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &Node) {
     use wacore::iq::tctoken::parse_privacy_token_notification;
     use wacore::store::traits::TcTokenEntry;
 
-    // Resolve the sender to a LID JID for storage.
+    let from_jid = node.attrs().optional_jid("from");
+
+    // Resolve the sender to a LID key for storage.
     // WA Web uses `sender_lid` attr if present, otherwise resolves from `from`.
     let sender_lid = node
         .attrs()
@@ -557,27 +559,27 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &Node) {
         Some(lid) if !lid.is_empty() => lid,
         _ => {
             // Fall back to resolving from the `from` JID via LID-PN cache
-            let from_jid = match node.attrs().optional_jid("from") {
+            let from = match &from_jid {
                 Some(jid) => jid,
                 None => {
                     warn!(target: "Client/TcToken", "privacy_token notification missing 'from' attribute");
                     return;
                 }
             };
 
-            if from_jid.is_lid() {
-                from_jid.user.clone()
+            if from.is_lid() {
+                from.user.clone()
             } else {
                 // Try to resolve phone number to LID
-                match client.lid_pn_cache.get_current_lid(&from_jid.user).await {
+                match client.lid_pn_cache.get_current_lid(&from.user).await {
                     Some(lid) => lid,
                     None => {
                         debug!(
                             target: "Client/TcToken",
                             "Cannot resolve LID for privacy_token sender {}, storing under PN",
-                            from_jid
+                            from
                         );
-                        from_jid.user.clone()
+                        from.user.clone()
                     }
                 }
             }
@@ -599,10 +601,26 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &Node) {
     }
 
     let backend = client.persistence_manager.backend();
+    let mut token_stored = false;
 
     for received in &received_tokens {
         match backend.get_tc_token(&sender_lid).await {
             Ok(Some(existing)) => {
+                // Skip if token bytes are identical and timestamp hasn't advanced
+                if existing.token == received.token {
+                    if received.timestamp > existing.token_timestamp {
+                        // Same bytes but newer timestamp — refresh to prevent premature pruning
+                        let refreshed = TcTokenEntry {
+                            token_timestamp: received.timestamp,
+                            ..existing
+                        };
+                        if let Err(e) = backend.put_tc_token(&sender_lid, &refreshed).await {
+                            warn!(target: "Client/TcToken", "Failed to refresh tc_token timestamp for {}: {e}", sender_lid);
+                        }
+                    }
+                    continue;
+                }
+
                 // Timestamp monotonicity guard: only store if incoming >= existing
                 if received.timestamp < existing.token_timestamp {
                     debug!(
@@ -624,6 +642,7 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &Node) {
                     warn!(target: "Client/TcToken", "Failed to update tc_token for {}: {e}", sender_lid);
                 } else {
                     debug!(target: "Client/TcToken", "Updated tc_token for {} (t={})", sender_lid, received.timestamp);
+                    token_stored = true;
                 }
             }
             Ok(None) => {
@@ -638,13 +657,22 @@ async fn handle_privacy_token_notification(client: &Arc<Client>, node: &Node) {
                     warn!(target: "Client/TcToken", "Failed to store tc_token for {}: {e}", sender_lid);
                 } else {
                     debug!(target: "Client/TcToken", "Stored new tc_token for {} (t={})", sender_lid, received.timestamp);
+                    token_stored = true;
                 }
             }
             Err(e) => {
                 warn!(target: "Client/TcToken", "Failed to read tc_token for {}: {e}, skipping", sender_lid);
             }
         }
     }
+
+    // Re-subscribe presence with the updated token.
+    if token_stored
+        && let Some(from) = &from_jid
+        && let Err(e) = client.presence().re_subscribe_when_active(from).await
+    {
+        debug!(target: "Client/TcToken", "Failed to re-subscribe presence for {from}: {e}");
+    }
 }
 
 /// Handle business notification (WhatsApp Web: `WAWebHandleBusinessNotification`).

tests/e2e/src/lib.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use wacore::store::traits::TcTokenEntry;
 use wacore::types::events::{Event, EventHandler};
 use whatsapp_rust::Jid;
 use whatsapp_rust::bot::Bot;
@@ -172,6 +173,46 @@ impl TestClient {
             .to_non_ad()
     }
 
+    /// Get the storage key used for this client's tcToken entries.
+    ///
+    /// Notification handling stores tcTokens under the sender's LID when it is
+    /// available, otherwise it falls back to the phone-number user part.
+    pub async fn tc_token_key(&self) -> anyhow::Result<String> {
+        if let Some(lid) = self.client.get_lid().await {
+            return Ok(lid.user);
+        }
+
+        self.client
+            .get_pn()
+            .await
+            .map(|jid| jid.user)
+            .ok_or_else(|| anyhow::anyhow!("Client should have a JID after connect"))
+    }
+
+    /// Wait until a tcToken entry exists for the given storage key.
+    pub async fn wait_for_tc_token(
+        &self,
+        jid_key: &str,
+        timeout_secs: u64,
+    ) -> anyhow::Result<TcTokenEntry> {
+        let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(timeout_secs);
+
+        loop {
+            if let Some(entry) = self.client.tc_token().get(jid_key).await? {
+                return Ok(entry);
+            }
+
+            if tokio::time::Instant::now() >= deadline {
+                return Err(anyhow::anyhow!(
+                    "Timed out waiting for tc_token entry for {}",
+                    jid_key
+                ));
+            }
+
+            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+        }
+    }
+
     // ── Event waiting ───────────────────────────────────────────────────────
 
     /// Wait for an event matching the predicate, with a timeout in seconds.