From e0bbf6b61f49f051fbc9f2fca9a82863b438c78c Mon Sep 17 00:00:00 2001
From: Jeff DeCew <jeffdq@google.com>
Date: Sat, 6 Feb 2021 09:39:59 -0500
Subject: [PATCH] Correct colors on forcibly decorated custom notifications

Test: manual
Fixes: 179415446
Change-Id: I9bebd087acfdcf4ffda8ff90deb8c1e806f51106
---
 ...otificationDecoratedCustomViewWrapper.java | 32 +++++++++++++++----
 .../row/wrapper/NotificationViewWrapper.java  |  3 ++
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationDecoratedCustomViewWrapper.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationDecoratedCustomViewWrapper.java
index 79648457c521c..301c3726793ae 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationDecoratedCustomViewWrapper.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationDecoratedCustomViewWrapper.java
@@ -29,6 +29,30 @@ public class NotificationDecoratedCustomViewWrapper extends NotificationTemplate
 
     private View mWrappedView = null;
 
+    /**
+     * Determines if the standard template contains a custom view, injected by Notification.Builder
+     */
+    public static boolean hasCustomView(View v) {
+        return getWrappedCustomView(v) != null;
+    }
+
+    private static View getWrappedCustomView(View view) {
+        if (view == null) {
+            return null;
+        }
+        ViewGroup container = view.findViewById(
+                com.android.internal.R.id.notification_main_column);
+        if (container == null) {
+            return null;
+        }
+        Integer childIndex = (Integer) container.getTag(
+                com.android.internal.R.id.notification_custom_view_index_tag);
+        if (childIndex == null || childIndex == -1) {
+            return null;
+        }
+        return container.getChildAt(childIndex);
+    }
+
     protected NotificationDecoratedCustomViewWrapper(Context ctx, View view,
             ExpandableNotificationRow row) {
         super(ctx, view, row);
@@ -36,13 +60,7 @@ public class NotificationDecoratedCustomViewWrapper extends NotificationTemplate
 
     @Override
     public void onContentUpdated(ExpandableNotificationRow row) {
-        ViewGroup container = mView.findViewById(
-                com.android.internal.R.id.notification_main_column);
-        Integer childIndex = (Integer) container.getTag(
-                com.android.internal.R.id.notification_custom_view_index_tag);
-        if (childIndex != null && childIndex != -1) {
-            mWrappedView = container.getChildAt(childIndex);
-        }
+        mWrappedView = getWrappedCustomView(mView);
 
         // Custom views will most likely use just white or black as their text color.
         // We need to scan through and replace these colors by Material NEXT colors.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
index 7b5c5f6dcf8f9..b3d1a94beaa95 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
@@ -80,6 +80,9 @@ public abstract class NotificationViewWrapper implements TransformableView {
             if (Notification.DecoratedCustomViewStyle.class.equals(style)) {
                 return new NotificationDecoratedCustomViewWrapper(ctx, v, row);
             }
+            if (NotificationDecoratedCustomViewWrapper.hasCustomView(v)) {
+                return new NotificationDecoratedCustomViewWrapper(ctx, v, row);
+            }
             return new NotificationTemplateViewWrapper(ctx, v, row);
         } else if (v instanceof NotificationHeaderView) {
             return new NotificationHeaderViewWrapper(ctx, v, row);