Add getAttributionSource function to mocked context

Bug: 221772127
Test: make, atest
Change-Id: I6671f27e5c600762b106fb79b1c79f7b47bc1a50

1 files changed, 11 insertions, 0 deletions

diff --git a/test-runner/src/android/test/IsolatedContext.java b/test-runner/src/android/test/IsolatedContext.java
index dd4a9a3a4d69..d5f92a3181b9 100644
--- a/test-runner/src/android/test/IsolatedContext.java
+++ b/test-runner/src/android/test/IsolatedContext.java
@@ -17,6 +17,7 @@
 package android.test;
 
 import android.accounts.AccountManager;
+import android.content.AttributionSource;
 import android.content.BroadcastReceiver;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -26,6 +27,7 @@ import android.content.IntentFilter;
 import android.content.ServiceConnection;
 import android.content.pm.PackageManager;
 import android.net.Uri;
+import android.os.Process;
 import android.test.mock.MockAccountManager;
 
 import java.io.File;
@@ -64,6 +66,15 @@ public class IsolatedContext extends ContextWrapper {
     }
 
     @Override
+    public AttributionSource getAttributionSource() {
+        AttributionSource attributionSource = super.getAttributionSource();
+        if (attributionSource == null) {
+            return new AttributionSource.Builder(Process.myUid()).build();
+        }
+        return attributionSource;
+    }
+
+    @Override
     public ContentResolver getContentResolver() {
         // We need to return the real resolver so that MailEngine.makeRight can get to the
         // subscribed feeds provider. TODO: mock out subscribed feeds too.