1 files changed, 51 insertions, 16 deletions

diff --git a/adservices/tests/unittest/service-core/src/com/android/adservices/service/measurement/reporting/AggregateReportBodyTest.java b/adservices/tests/unittest/service-core/src/com/android/adservices/service/measurement/reporting/AggregateReportBodyTest.java
index 3d79b2e089..d285a741ae 100644
--- a/adservices/tests/unittest/service-core/src/com/android/adservices/service/measurement/reporting/AggregateReportBodyTest.java
+++ b/adservices/tests/unittest/service-core/src/com/android/adservices/service/measurement/reporting/AggregateReportBodyTest.java
@@ -22,8 +22,10 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import com.android.adservices.HpkeJni;
+import com.android.adservices.service.measurement.aggregation.AggregateCryptoConverter;
 import com.android.adservices.service.measurement.aggregation.AggregateCryptoFixture;
 import com.android.adservices.service.measurement.aggregation.AggregateEncryptionKey;
+import com.android.adservices.service.measurement.util.UnsignedLong;
 
 import org.json.JSONArray;
 import org.json.JSONException;
@@ -50,8 +52,8 @@ public class AggregateReportBodyTest {
     private static final String SCHEDULED_REPORT_TIME = "1246174158155";
     private static final String VERSION = "12";
     private static final String REPORT_ID = "A1";
-    private static final Long SOURCE_DEBUG_KEY = 27628792L;
-    private static final Long TRIGGER_DEBUG_KEY = 23443234L;
+    private static final UnsignedLong SOURCE_DEBUG_KEY = new UnsignedLong(27628792L);
+    private static final UnsignedLong TRIGGER_DEBUG_KEY = new UnsignedLong(23443234L);
     private static final String REPORTING_ORIGIN = "https://adtech.domain";
     private static final String DEBUG_CLEARTEXT_PAYLOAD = "{\"operation\":\"histogram\","
             + "\"data\":[{\"bucket\":1369,\"value\":32768},{\"bucket\":3461,"
@@ -123,10 +125,8 @@ public class AggregateReportBodyTest {
         assertEquals(REPORTING_ORIGIN, sharedInfoJson.get("reporting_origin"));
         assertEquals(ATTRIBUTION_DESTINATION, sharedInfoJson.get("attribution_destination"));
         assertEquals(SOURCE_REGISTRATION_TIME, sharedInfoJson.get("source_registration_time"));
-        assertEquals(
-                Long.toUnsignedString(SOURCE_DEBUG_KEY), aggregateJson.get("source_debug_key"));
-        assertEquals(
-                Long.toUnsignedString(TRIGGER_DEBUG_KEY), aggregateJson.get("trigger_debug_key"));
+        assertEquals(SOURCE_DEBUG_KEY.toString(), aggregateJson.get("source_debug_key"));
+        assertEquals(TRIGGER_DEBUG_KEY.toString(), aggregateJson.get("trigger_debug_key"));
     }
 
     @Test
@@ -158,8 +158,7 @@ public class AggregateReportBodyTest {
         assertEquals(REPORTING_ORIGIN, sharedInfoJson.get("reporting_origin"));
         assertEquals(ATTRIBUTION_DESTINATION, sharedInfoJson.get("attribution_destination"));
         assertEquals(SOURCE_REGISTRATION_TIME, sharedInfoJson.get("source_registration_time"));
-        assertEquals(
-                Long.toUnsignedString(SOURCE_DEBUG_KEY), aggregateJson.get("source_debug_key"));
+        assertEquals(SOURCE_DEBUG_KEY.toString(), aggregateJson.get("source_debug_key"));
         assertNull(aggregateJson.opt("trigger_debug_key"));
     }
 
@@ -177,8 +176,7 @@ public class AggregateReportBodyTest {
         assertEquals(ATTRIBUTION_DESTINATION, sharedInfoJson.get("attribution_destination"));
         assertEquals(SOURCE_REGISTRATION_TIME, sharedInfoJson.get("source_registration_time"));
         assertNull(aggregateJson.opt("source_debug_key"));
-        assertEquals(
-                Long.toUnsignedString(TRIGGER_DEBUG_KEY), aggregateJson.get("trigger_debug_key"));
+        assertEquals(TRIGGER_DEBUG_KEY.toString(), aggregateJson.get("trigger_debug_key"));
     }
 
     @Test
@@ -196,13 +194,50 @@ public class AggregateReportBodyTest {
         assertEncryptedPayload(aggregateServicePayloads);
     }
 
-    private void assertEncodedDebugPayload(JSONObject aggregateServicePayloads) throws Exception {
-        final String encodedPayloadBase64 =
-                (String) aggregateServicePayloads.get("debug_cleartext_payload");
-        assertNotNull(encodedPayloadBase64);
+    @Test
+    public void testAggregationServicePayloadsJsonSerializationWithDebugKey() throws Exception {
+        AggregateReportBody aggregateReport =
+                createAggregateReportBodyExampleWithSingleTriggerDebugKey();
+
+        AggregateEncryptionKey key = AggregateCryptoFixture.getKey();
+        JSONArray aggregationServicePayloadsJson =
+                aggregateReport.aggregationServicePayloadsToJson(/* sharedInfo = */ null, key);
+
+        JSONObject aggregateServicePayloads = aggregationServicePayloadsJson.getJSONObject(0);
 
-        final byte[] cborEncodedPayload = Base64.getDecoder().decode(encodedPayloadBase64);
-        assertCborEncoded(cborEncodedPayload);
+        assertEquals(key.getKeyId(), aggregateServicePayloads.get("key_id"));
+        assertEquals(
+                AggregateCryptoConverter.encode(DEBUG_CLEARTEXT_PAYLOAD),
+                aggregateServicePayloads.opt("debug_cleartext_payload"));
+        assertEncodedDebugPayload(aggregateServicePayloads);
+        assertEncryptedPayload(aggregateServicePayloads);
+    }
+
+    @Test
+    public void testAggregationServicePayloadsJsonSerializationWithoutDebugKey() throws Exception {
+        AggregateReportBody aggregateReport = createAggregateReportBodyExampleWithNullDebugKeys();
+
+        AggregateEncryptionKey key = AggregateCryptoFixture.getKey();
+        JSONArray aggregationServicePayloadsJson =
+                aggregateReport.aggregationServicePayloadsToJson(/* sharedInfo = */ null, key);
+
+        JSONObject aggregateServicePayloads = aggregationServicePayloadsJson.getJSONObject(0);
+
+        assertEquals(key.getKeyId(), aggregateServicePayloads.get("key_id"));
+        assertNull(aggregateServicePayloads.opt("debug_cleartext_payload"));
+        assertEncodedDebugPayload(aggregateServicePayloads);
+        assertEncryptedPayload(aggregateServicePayloads);
+    }
+
+    private void assertEncodedDebugPayload(JSONObject aggregateServicePayloads) throws Exception {
+        if (!aggregateServicePayloads.isNull("debug_cleartext_payload")) {
+            final String encodedPayloadBase64 =
+                    (String) aggregateServicePayloads.get("debug_cleartext_payload");
+            assertNotNull(encodedPayloadBase64);
+
+            final byte[] cborEncodedPayload = Base64.getDecoder().decode(encodedPayloadBase64);
+            assertCborEncoded(cborEncodedPayload);
+        }
     }
 
     private void assertEncryptedPayload(JSONObject aggregateServicePayloads) throws Exception {