Stop CoreData from skipping TES on NightOwl replies and retweets

home-mixer/candidate_hydrators/core_data_candidate_hydrator.rs

@@ -38,7 +38,19 @@ impl CachedHydrator<ScoredPostsQuery, PostCandidate> for CoreDataCandidateHydrat
     }
 
     fn already_hydrated(&self, candidate: &PostCandidate) -> bool {
-        candidate.author_id != 0 && !candidate.tweet_text.is_empty()
+        if candidate.author_id == 0 || candidate.tweet_text.is_empty() {
+            return false;
+        }
+        // NightOwl (Latest Following) pre-fills author_id and tweet_text.
+        // Replies still need TES for ancestor_users (SelfReplyChainFilter).
+        // Retweets still need TES for retweeted_user_id.
+        if candidate.in_reply_to_tweet_id.is_some() && candidate.ancestor_users.is_empty() {
+            return false;
+        }
+        if candidate.retweeted_tweet_id.is_some() && candidate.retweeted_user_id.is_none() {
+            return false;
+        }
+        true
     }
 
     fn cache_store(&self) -> &dyn CacheStore<Self::CacheKey, Self::CacheValue> {
@@ -104,7 +116,10 @@ impl CachedHydrator<ScoredPostsQuery, PostCandidate> for CoreDataCandidateHydrat
                 }
                 Some(Ok(None)) | None => {
                     missing_count += 1;
-                    hydrated_candidates.push(Ok(PostCandidate::default()));
+                    // Keep source-populated fields. An empty default would
+                    // clear NightOwl reply/retweet ids now that we call TES
+                    // for those cards.
+                    hydrated_candidates.push(Ok(source_core_fields(candidate)));
                 }
                 Some(Err(err)) => {
                     hydrated_candidates.push(Err(err.to_string()));
@@ -129,6 +144,18 @@ impl CachedHydrator<ScoredPostsQuery, PostCandidate> for CoreDataCandidateHydrat
     }
 }
 
+fn source_core_fields(candidate: &PostCandidate) -> PostCandidate {
+    PostCandidate {
+        author_id: candidate.author_id,
+        retweeted_user_id: candidate.retweeted_user_id,
+        retweeted_tweet_id: candidate.retweeted_tweet_id,
+        in_reply_to_tweet_id: candidate.in_reply_to_tweet_id,
+        ancestor_users: candidate.ancestor_users.clone(),
+        tweet_text: candidate.tweet_text.clone(),
+        ..Default::default()
+    }
+}
+
 fn core_data_fetch_ids(candidates: &[PostCandidate]) -> Vec<u64> {
     let mut fetch_ids: Vec<u64> = candidates.iter().map(|c| c.tweet_id).collect();
     fetch_ids.extend(
@@ -178,3 +205,176 @@ impl CoreDataCandidateHydrator {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::clients::tweet_entity_service_client::MockTESClient;
+    use crate::filters::self_reply_chain_filter::SelfReplyChainFilter;
+    use crate::models::user_features::UserFeatures;
+    use xai_candidate_pipeline::filter::Filter;
+    use xai_candidate_pipeline::hydrator::Hydrator;
+
+    fn nightowl_reply() -> PostCandidate {
+        PostCandidate {
+            tweet_id: 30,
+            author_id: 10,
+            in_reply_to_tweet_id: Some(20),
+            ancestors: vec![20, 50],
+            tweet_text: "self reply in someone else's thread".into(),
+            ..Default::default()
+        }
+    }
+
+    fn nightowl_retweet() -> PostCandidate {
+        PostCandidate {
+            tweet_id: 31,
+            author_id: 10,
+            retweeted_tweet_id: Some(21),
+            tweet_text: "rt".into(),
+            ..Default::default()
+        }
+    }
+
+    fn nightowl_original() -> PostCandidate {
+        PostCandidate {
+            tweet_id: 32,
+            author_id: 10,
+            tweet_text: "hello".into(),
+            ..Default::default()
+        }
+    }
+
+    fn tes_reply() -> PureCoreData {
+        PureCoreData {
+            author_id: 10,
+            text: "self reply in someone else's thread".into(),
+            in_reply_to_tweet_id: Some(20),
+            in_reply_to_user_id: Some(10),
+            ..Default::default()
+        }
+    }
+
+    fn tes_root() -> PureCoreData {
+        PureCoreData {
+            author_id: 99,
+            text: "root from unfollowed user".into(),
+            ..Default::default()
+        }
+    }
+
+    fn tes_retweet() -> PureCoreData {
+        PureCoreData {
+            author_id: 10,
+            text: "rt".into(),
+            source_tweet_id: Some(21),
+            source_user_id: Some(77),
+            ..Default::default()
+        }
+    }
+
+    async fn hydrator(core_data: HashMap<u64, Option<PureCoreData>>) -> CoreDataCandidateHydrator {
+        CoreDataCandidateHydrator::new(Arc::new(MockTESClient {
+            core_data,
+            ..Default::default()
+        }) as Arc<dyn TESClient + Send + Sync>)
+        .await
+    }
+
+    async fn run(
+        hydrator: &CoreDataCandidateHydrator,
+        candidates: Vec<PostCandidate>,
+    ) -> Vec<PostCandidate> {
+        let mut candidates = candidates;
+        let hydrated = Hydrator::hydrate(hydrator, &ScoredPostsQuery::default(), &candidates).await;
+        hydrator.update_all(&mut candidates, hydrated);
+        candidates
+    }
+
+    #[tokio::test]
+    async fn already_hydrated_skips_complete_originals() {
+        let h = hydrator(HashMap::new()).await;
+        assert!(h.already_hydrated(&nightowl_original()));
+        assert!(!h.already_hydrated(&nightowl_reply()));
+        assert!(!h.already_hydrated(&nightowl_retweet()));
+        assert!(!h.already_hydrated(&PostCandidate {
+            author_id: 10,
+            tweet_text: String::new(),
+            ..Default::default()
+        }));
+    }
+
+    #[tokio::test]
+    async fn nightowl_reply_gets_ancestor_users_from_tes() {
+        let mut core_data = HashMap::new();
+        core_data.insert(30, Some(tes_reply()));
+        core_data.insert(50, Some(tes_root()));
+        let h = hydrator(core_data).await;
+
+        let out = run(&h, vec![nightowl_reply()]).await;
+        assert_eq!(out[0].ancestor_users, vec![10, 99]);
+        assert_eq!(out[0].in_reply_to_tweet_id, Some(20));
+        assert_eq!(out[0].author_id, 10);
+        assert_eq!(out[0].tweet_text, "self reply in someone else's thread");
+    }
+
+    #[tokio::test]
+    async fn nightowl_retweet_gets_original_author_from_tes() {
+        let mut core_data = HashMap::new();
+        core_data.insert(31, Some(tes_retweet()));
+        let h = hydrator(core_data).await;
+
+        let out = run(&h, vec![nightowl_retweet()]).await;
+        assert_eq!(out[0].retweeted_tweet_id, Some(21));
+        assert_eq!(out[0].retweeted_user_id, Some(77));
+    }
+
+    #[tokio::test]
+    async fn nightowl_original_still_skips_tes() {
+        let mut core_data = HashMap::new();
+        core_data.insert(
+            32,
+            Some(PureCoreData {
+                author_id: 10,
+                text: "tes text must not replace nightowl".into(),
+                ..Default::default()
+            }),
+        );
+        let h = hydrator(core_data).await;
+        let out = run(&h, vec![nightowl_original()]).await;
+        assert_eq!(out[0].tweet_text, "hello");
+    }
+
+    #[tokio::test]
+    async fn tes_miss_keeps_nightowl_reply_ids() {
+        let h = hydrator(HashMap::new()).await;
+        let out = run(&h, vec![nightowl_reply()]).await;
+        assert_eq!(out[0].in_reply_to_tweet_id, Some(20));
+        assert_eq!(out[0].ancestors, vec![20, 50]);
+        assert_eq!(out[0].author_id, 10);
+        assert_eq!(out[0].tweet_text, "self reply in someone else's thread");
+        assert!(out[0].ancestor_users.is_empty());
+    }
+
+    #[tokio::test]
+    async fn self_reply_chain_drops_after_nightowl_tes_hydrate() {
+        let mut core_data = HashMap::new();
+        core_data.insert(30, Some(tes_reply()));
+        core_data.insert(50, Some(tes_root()));
+        let h = hydrator(core_data).await;
+        let out = run(&h, vec![nightowl_reply()]).await;
+
+        let query = ScoredPostsQuery {
+            user_id: 1,
+            user_features: UserFeatures {
+                followed_user_ids: vec![10],
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+        let result = SelfReplyChainFilter.filter(&query, out);
+        assert_eq!(result.removed.len(), 1);
+        assert_eq!(result.removed[0].tweet_id, 30);
+        assert!(result.kept.is_empty());
+    }
+}