Drop retweets when the original author blocked the viewer

home-mixer/candidate_hydrators/blocked_by_hydrator.rs

@@ -1,5 +1,6 @@
 use crate::models::candidate::PostCandidate;
 use crate::models::query::ScoredPostsQuery;
+use std::collections::HashSet;
 use std::sync::Arc;
 use tonic::async_trait;
 use xai_candidate_pipeline::component_library::clients::SocialGraphClientOps;
@@ -26,11 +27,16 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for BlockedByHydrator {
         query: &ScoredPostsQuery,
         candidates: &[PostCandidate],
     ) -> Vec<Result<PostCandidate, String>> {
-        let author_ids: Vec<u64> = candidates.iter().map(|x| x.author_id).collect();
+        let user_ids: Vec<u64> = candidates
+            .iter()
+            .flat_map(|c| std::iter::once(c.author_id).chain(c.retweeted_user_id))
+            .collect::<HashSet<_>>()
+            .into_iter()
+            .collect();
 
         let blocked_by_user_ids = match self
             .socialgraph_client
-            .check_blocked_by(query.user_id, &author_ids)
+            .check_blocked_by(query.user_id, &user_ids)
             .await
         {
             Ok(ids) => ids,
@@ -42,7 +48,10 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for BlockedByHydrator {
         candidates
             .iter()
             .map(|candidate| {
-                let author_blocks_viewer = blocked_by_user_ids.contains(&candidate.author_id);
+                let author_blocks_viewer = blocked_by_user_ids.contains(&candidate.author_id)
+                    || candidate
+                        .retweeted_user_id
+                        .is_some_and(|uid| blocked_by_user_ids.contains(&uid));
                 Ok(PostCandidate {
                     author_blocks_viewer: Some(author_blocks_viewer),
                     ..Default::default()
@@ -55,3 +64,112 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for BlockedByHydrator {
         candidate.author_blocks_viewer = hydrated.author_blocks_viewer;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tonic::Status;
+
+    struct MockSocialGraph {
+        blocked_by: HashSet<u64>,
+    }
+
+    #[async_trait]
+    impl SocialGraphClientOps for MockSocialGraph {
+        async fn get_following_list(&self, _user_id: u64) -> Result<Vec<u64>, Status> {
+            Ok(vec![])
+        }
+        async fn check_blocked_by(
+            &self,
+            _viewer_id: u64,
+            author_ids: &[u64],
+        ) -> Result<HashSet<u64>, Status> {
+            Ok(author_ids
+                .iter()
+                .copied()
+                .filter(|id| self.blocked_by.contains(id))
+                .collect())
+        }
+        async fn check_followed_by(
+            &self,
+            _viewer_id: u64,
+            _user_ids: &[u64],
+        ) -> Result<HashSet<u64>, Status> {
+            Ok(HashSet::new())
+        }
+        async fn get_blocked_user_ids(&self, _viewer_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_muted_user_ids(&self, _viewer_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_followed_user_ids(&self, _viewer_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_follower_ids(&self, _user_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_subscribed_user_ids(&self, _viewer_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_device_following_user_ids(&self, _viewer_id: u64) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+        async fn get_hide_recommendations_user_ids(
+            &self,
+            _viewer_id: u64,
+        ) -> Result<Vec<i64>, Status> {
+            Ok(vec![])
+        }
+    }
+
+    fn hydrator(blocked_by: HashSet<u64>) -> BlockedByHydrator {
+        BlockedByHydrator {
+            socialgraph_client: Arc::new(MockSocialGraph { blocked_by }),
+        }
+    }
+
+    #[tokio::test]
+    async fn retweet_of_author_who_blocked_viewer_is_marked() {
+        let hydrator = hydrator(HashSet::from([99]));
+        let rt = PostCandidate {
+            tweet_id: 1,
+            author_id: 10,
+            retweeted_user_id: Some(99),
+            ..Default::default()
+        };
+        let native = PostCandidate {
+            tweet_id: 2,
+            author_id: 10,
+            ..Default::default()
+        };
+        let query = ScoredPostsQuery {
+            user_id: 1,
+            ..Default::default()
+        };
+        let hydrated = hydrator.hydrate(&query, &[rt.clone(), native.clone()]).await;
+        let mut rt = rt;
+        let mut native = native;
+        hydrator.update(&mut rt, hydrated[0].clone().unwrap());
+        hydrator.update(&mut native, hydrated[1].clone().unwrap());
+        assert_eq!(rt.author_blocks_viewer, Some(true));
+        assert_eq!(native.author_blocks_viewer, Some(false));
+    }
+
+    #[tokio::test]
+    async fn primary_author_who_blocked_viewer_is_still_marked() {
+        let hydrator = hydrator(HashSet::from([10]));
+        let mut candidate = PostCandidate {
+            tweet_id: 1,
+            author_id: 10,
+            ..Default::default()
+        };
+        let query = ScoredPostsQuery {
+            user_id: 1,
+            ..Default::default()
+        };
+        let hydrated = hydrator.hydrate(&query, &[candidate.clone()]).await;
+        hydrator.update(&mut candidate, hydrated[0].clone().unwrap());
+        assert_eq!(candidate.author_blocks_viewer, Some(true));
+    }
+}