From 7974e19b7d20d8e1ef0261df9c3a41dedd6021b8 Mon Sep 17 00:00:00 2001
From: Rafael Higuera Silva <rafahs@google.com>
Date: Mon, 28 Feb 2022 19:16:34 +0000
Subject: [PATCH] Add getAttributionSource function to mocked context

Bug: 221772127
Test: make, atest
Change-Id: I6671f27e5c600762b106fb79b1c79f7b47bc1a50
---
 test-runner/src/android/test/IsolatedContext.java | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/test-runner/src/android/test/IsolatedContext.java b/test-runner/src/android/test/IsolatedContext.java
index dd4a9a3a4d697..d5f92a3181b93 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;
@@ -63,6 +65,15 @@ public class IsolatedContext extends ContextWrapper {
         return intents;
     }
 
+    @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