Add roundtripping of cap_key in sposored content impression callback urls

Author: jonesetc

components/ads-client/src/client.rs

@@ -121,7 +121,7 @@ where
 
     pub fn record_impression(
         &self,
-        impression_url: Url,
+        mut impression_url: Url,
         ohttp: bool,
     ) -> Result<(), RecordImpressionError> {
         // TODO: Re-enable cache invalidation behind a Nimbus experiment.
@@ -130,6 +130,26 @@ where
         // if let Some(request_hash) = pop_request_hash_from_url(&mut impression_url) {
         //     let _ = self.client.invalidate_cache_by_hash(&request_hash);
         // }
+
+        // TODO: Add count call with _cap_key for impression capping logic
+        if let Some((_, _cap_key)) = impression_url
+            .query_pairs()
+            .find(|(key, _)| key == "cap_key")
+        {
+            let original_url = impression_url.clone();
+            impression_url
+                .query_pairs_mut()
+                .clear()
+                .extend_pairs(
+                    original_url
+                        .query_pairs()
+                        .collect::<Vec<_>>()
+                        .iter()
+                        .filter(|(key, _)| key != "cap_key"),
+                )
+                .finish();
+        }
+
         self.client
             .record_impression(impression_url, ohttp)
             .inspect_err(|e| {
@@ -368,6 +388,33 @@ mod tests {
         m.assert();
     }
 
+    #[test]
+    fn test_record_impression_removes_cap_key() {
+        viaduct_dev::init_backend_dev();
+        let mars_client = MARSClient::new(Environment::Test, None, MozAdsTelemetryWrapper::noop());
+        let ads_client = new_with_mars_client(mars_client);
+
+        let base_url = mockito::server_url();
+        let path_and_query = "/impression?kept=example";
+        let callback_url = Url::parse(&format!("{}{}", base_url, path_and_query)).unwrap();
+
+        let mock = mockito::mock("GET", path_and_query)
+            .with_status(200)
+            .create();
+
+        ads_client.record_impression(callback_url, false).unwrap();
+
+        mock.assert();
+
+        let callback_url_with_cap_key =
+            Url::parse(&format!("{}{}&cap_key=test", base_url, path_and_query)).unwrap();
+        ads_client
+            .record_impression(callback_url_with_cap_key, false)
+            .unwrap();
+
+        mock.expect(2).assert();
+    }
+
     #[test]
     #[ignore = "Cache invalidation temporarily disabled - will be re-enabled behind Nimbus experiment"]
     fn test_record_click_invalidates_cache() {

components/ads-client/src/mars/ad_response.rs

@@ -47,6 +47,7 @@ impl<A: AdResponseValue> AdResponse<A> {
         let hash_str = request_hash.to_string();
         for (placement_id, ads) in self.data.iter_mut() {
             for (position, ad) in ads.iter_mut().enumerate() {
+                let impression_enriching_params = ad.impression_enriching_params();
                 let callbacks = ad.callbacks_mut();
                 callbacks
                     .click
@@ -55,7 +56,8 @@ impl<A: AdResponseValue> AdResponse<A> {
                 callbacks
                     .impression
                     .query_pairs_mut()
-                    .append_pair("request_hash", &hash_str);
+                    .append_pair("request_hash", &hash_str)
+                    .extend_pairs(impression_enriching_params);
                 if let Some(report_url) = callbacks.report.as_mut() {
                     report_url
                         .query_pairs_mut()
@@ -161,6 +163,9 @@ pub struct AdCallbacks {
 
 pub trait AdResponseValue: DeserializeOwned {
     fn callbacks_mut(&mut self) -> &mut AdCallbacks;
+    fn impression_enriching_params(&self) -> Vec<(String, String)> {
+        vec![]
+    }
 }
 
 impl AdResponseValue for AdImage {
@@ -173,6 +178,10 @@ impl AdResponseValue for AdSpoc {
     fn callbacks_mut(&mut self) -> &mut AdCallbacks {
         &mut self.callbacks
     }
+
+    fn impression_enriching_params(&self) -> Vec<(String, String)> {
+        vec![("cap_key".to_owned(), self.caps.cap_key.clone())]
+    }
 }
 
 impl AdResponseValue for AdTile {
@@ -514,6 +523,86 @@ mod tests {
         assert!(report_1.contains("position=1"));
     }
 
+    #[test]
+    fn test_enrich_callbacks_spoc_impressions_have_cap_key() {
+        let mut response = AdResponse {
+            data: HashMap::from([(
+                "tile1".into(),
+                vec![
+                    AdSpoc {
+                        block_key: "block_key1".into(),
+                        callbacks: AdCallbacks {
+                            click: Url::parse("https://example.com/click1").unwrap(),
+                            impression: Url::parse("https://example.com/impression1").unwrap(),
+                            report: Some(Url::parse("https://example.com/report1").unwrap()),
+                        },
+                        caps: SpocFrequencyCaps {
+                            cap_key: "cap_key1".into(),
+                            day: 100,
+                        },
+                        domain: "1.example.com".into(),
+                        excerpt: "excerpt1".into(),
+                        format: "format1".into(),
+                        image_url: "https://example.com/image1.png".into(),
+                        url: "https://example.com/ad1".into(),
+                        ranking: SpocRanking {
+                            priority: 1,
+                            personalization_models: None,
+                            item_score: 1.0,
+                        },
+                        sponsor: "sponsor1".into(),
+                        sponsored_by_override: None,
+                        title: "title1".into(),
+                    },
+                    AdSpoc {
+                        block_key: "block_key2".into(),
+                        callbacks: AdCallbacks {
+                            click: Url::parse("https://example.com/click2").unwrap(),
+                            impression: Url::parse("https://example.com/impression2").unwrap(),
+                            report: Some(Url::parse("https://example.com/report2").unwrap()),
+                        },
+                        caps: SpocFrequencyCaps {
+                            cap_key: "cap_key2".into(),
+                            day: 200,
+                        },
+                        domain: "2.example.com".into(),
+                        excerpt: "excerpt2".into(),
+                        format: "format2".into(),
+                        image_url: "https://example.com/image2.png".into(),
+                        url: "https://example.com/ad2".into(),
+                        ranking: SpocRanking {
+                            priority: 2,
+                            personalization_models: None,
+                            item_score: 2.0,
+                        },
+                        sponsor: "sponsor2".into(),
+                        sponsored_by_override: None,
+                        title: "title2".into(),
+                    },
+                ],
+            )]),
+        };
+
+        let request_hash = RequestHash::from("abc123def456");
+        response.enrich_callbacks(&request_hash);
+
+        let ads = &response.data["tile1"];
+
+        assert!(ads[0]
+            .callbacks
+            .impression
+            .query()
+            .unwrap_or("")
+            .contains("cap_key=cap_key1"));
+
+        assert!(ads[1]
+            .callbacks
+            .impression
+            .query()
+            .unwrap_or("")
+            .contains("cap_key=cap_key2"));
+    }
+
     #[test]
     fn test_enrich_callbacks_skips_ads_without_report_url() {
         let mut response = AdResponse {