1 files changed, 21 insertions, 5 deletions

diff --git a/adservices/tests/unittest/service-core/src/com/android/adservices/data/measurement/DbState.java b/adservices/tests/unittest/service-core/src/com/android/adservices/data/measurement/DbState.java
index e649226bc2..512374fb5f 100644
--- a/adservices/tests/unittest/service-core/src/com/android/adservices/data/measurement/DbState.java
+++ b/adservices/tests/unittest/service-core/src/com/android/adservices/data/measurement/DbState.java
@@ -26,6 +26,7 @@ import com.android.adservices.service.measurement.Source;
 import com.android.adservices.service.measurement.Trigger;
 import com.android.adservices.service.measurement.aggregation.AggregateEncryptionKey;
 import com.android.adservices.service.measurement.aggregation.AggregateReport;
+import com.android.adservices.service.measurement.util.UnsignedLong;
 
 import org.json.JSONArray;
 import org.json.JSONException;
@@ -219,13 +220,13 @@ public class DbState {
     private Source getSourceFrom(JSONObject sJSON) throws JSONException {
         return new Source.Builder()
                 .setId(sJSON.getString("id"))
-                .setEventId(sJSON.getLong("eventId"))
+                .setEventId(new UnsignedLong(sJSON.getString("eventId")))
                 .setSourceType(
                         Source.SourceType.valueOf(
                                 sJSON.getString("sourceType").toUpperCase(Locale.ENGLISH)))
                 .setPublisher(Uri.parse(sJSON.getString("publisher")))
                 .setPublisherType(sJSON.optInt("publisherType"))
-                .setAppDestination(Uri.parse(sJSON.getString("appDestination")))
+                .setAppDestination(parseIfNonNull(sJSON.optString("appDestination", null)))
                 .setWebDestination(parseIfNonNull(sJSON.optString("webDestination", null)))
                 .setAggregateSource(sJSON.optString("aggregationKeys", null))
                 .setAggregateContributions(sJSON.optInt("aggregateContributions"))
@@ -241,7 +242,7 @@ public class DbState {
                 .setInstallAttributed(sJSON.optBoolean("installAttributed", false))
                 .setAttributionMode(
                         sJSON.optInt("attribution_mode", Source.AttributionMode.TRUTHFULLY))
-                .setAggregateFilterData(sJSON.optString("aggregateFilterData", null))
+                .setFilterData(sJSON.optString("filterData", null))
                 .build();
     }
 
@@ -258,16 +259,17 @@ public class DbState {
                 .setStatus(tJSON.getInt("status"))
                 .setRegistrant(Uri.parse(tJSON.getString("registrant")))
                 .setFilters(tJSON.optString("filters", null))
+                .setNotFilters(tJSON.optString("not_filters", null))
                 .build();
     }
 
     private EventReport getEventReportFrom(JSONObject rJSON) throws JSONException {
         return new EventReport.Builder()
                 .setId(rJSON.getString("id"))
-                .setSourceId(rJSON.getLong("sourceId"))
+                .setSourceEventId(new UnsignedLong(rJSON.getString("sourceEventId")))
                 .setAttributionDestination(Uri.parse(rJSON.getString("attributionDestination")))
                 .setEnrollmentId(rJSON.getString("enrollmentId"))
-                .setTriggerData(rJSON.getLong("triggerData"))
+                .setTriggerData(new UnsignedLong(rJSON.getString("triggerData")))
                 .setTriggerTime(rJSON.getLong("triggerTime"))
                 .setReportTime(rJSON.getLong("reportTime"))
                 .setTriggerPriority(rJSON.getLong("triggerPriority"))
@@ -276,6 +278,8 @@ public class DbState {
                 .setSourceType(
                         Source.SourceType.valueOf(
                                 rJSON.getString("sourceType").toUpperCase(Locale.ENGLISH)))
+                .setSourceId(rJSON.optString("sourceId", null))
+                .setTriggerId(rJSON.optString("triggerId", null))
                 .build();
     }
 
@@ -290,6 +294,8 @@ public class DbState {
                 .setEnrollmentId(attrJSON.getString("enrollmentId"))
                 .setTriggerTime(attrJSON.getLong("triggerTime"))
                 .setRegistrant(attrJSON.getString("registrant"))
+                .setSourceId(attrJSON.optString("sourceId", null))
+                .setTriggerId(attrJSON.optString("triggerId", null))
                 .build();
     }
 
@@ -336,6 +342,14 @@ public class DbState {
                         cursor.getString(
                                 cursor.getColumnIndex(
                                         MeasurementTables.AttributionContract.REGISTRANT)))
+                .setSourceId(
+                        cursor.getString(
+                                cursor.getColumnIndex(
+                                        MeasurementTables.AttributionContract.SOURCE_ID)))
+                .setTriggerId(
+                        cursor.getString(
+                                cursor.getColumnIndex(
+                                        MeasurementTables.AttributionContract.TRIGGER_ID)))
                 .build();
     }
 
@@ -351,6 +365,8 @@ public class DbState {
                 .setDebugCleartextPayload(rJSON.getString("debugCleartextPayload"))
                 .setStatus(rJSON.getInt("status"))
                 .setApiVersion(rJSON.optString("apiVersion", null))
+                .setSourceId(rJSON.optString("sourceId", null))
+                .setTriggerId(rJSON.optString("triggerId", null))
                 .build();
     }