diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/TransitionAnimationHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/TransitionAnimationHelper.java
index 8d29901c3a07e..bcc37baa5b00c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/TransitionAnimationHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/TransitionAnimationHelper.java
@@ -143,8 +143,12 @@ public class TransitionAnimationHelper {
         Animation a = null;
         if (animAttr != 0) {
             if (overrideType == ANIM_FROM_STYLE && !isTask) {
-                a = loadCustomActivityTransition(animAttr, options, enter, transitionAnimation);
-                if (a == null) {
+                final TransitionInfo.AnimationOptions.CustomActivityTransition customTransition =
+                        getCustomActivityTransition(animAttr, options);
+                if (customTransition != null) {
+                    a = loadCustomActivityTransition(
+                            customTransition, options, enter, transitionAnimation);
+                } else {
                     a = transitionAnimation
                             .loadAnimationAttr(options.getPackageName(), options.getAnimations(),
                                     animAttr, translucent);
@@ -161,10 +165,8 @@ public class TransitionAnimationHelper {
         return a;
     }
 
-    static Animation loadCustomActivityTransition(int animAttr,
-            TransitionInfo.AnimationOptions options, boolean enter,
-            TransitionAnimation transitionAnimation) {
-        Animation a = null;
+    static TransitionInfo.AnimationOptions.CustomActivityTransition getCustomActivityTransition(
+            int animAttr, TransitionInfo.AnimationOptions options) {
         boolean isOpen = false;
         switch (animAttr) {
             case R.styleable.WindowAnimation_activityOpenEnterAnimation:
@@ -178,17 +180,19 @@ public class TransitionAnimationHelper {
                 return null;
         }
 
-        final TransitionInfo.AnimationOptions.CustomActivityTransition transitionAnim =
-                options.getCustomActivityTransition(isOpen);
-        if (transitionAnim != null) {
-            a = transitionAnimation.loadAppTransitionAnimation(options.getPackageName(),
-                    enter ? transitionAnim.getCustomEnterResId()
-                            : transitionAnim.getCustomExitResId());
-            if (a != null && transitionAnim.getCustomBackgroundColor() != 0) {
-                a.setBackdropColor(transitionAnim.getCustomBackgroundColor());
-            }
-        }
+        return options.getCustomActivityTransition(isOpen);
+    }
 
+    static Animation loadCustomActivityTransition(
+            @NonNull TransitionInfo.AnimationOptions.CustomActivityTransition transitionAnim,
+            TransitionInfo.AnimationOptions options, boolean enter,
+            TransitionAnimation transitionAnimation) {
+        final Animation a = transitionAnimation.loadAppTransitionAnimation(options.getPackageName(),
+                enter ? transitionAnim.getCustomEnterResId()
+                        : transitionAnim.getCustomExitResId());
+        if (a != null && transitionAnim.getCustomBackgroundColor() != 0) {
+            a.setBackdropColor(transitionAnim.getCustomBackgroundColor());
+        }
         return a;
     }
 
diff --git a/services/core/java/com/android/server/wm/AppTransition.java b/services/core/java/com/android/server/wm/AppTransition.java
index f73c68a42ec75..939cf1ae471be 100644
--- a/services/core/java/com/android/server/wm/AppTransition.java
+++ b/services/core/java/com/android/server/wm/AppTransition.java
@@ -888,8 +888,11 @@ public class AppTransition implements Dump {
         } else {
             int animAttr = mapOpenCloseTransitTypes(transit, enter);
             if (animAttr != 0) {
-                a = loadCustomActivityAnimation(animAttr, enter, container);
-                if (a == null) {
+                final CustomAppTransition customAppTransition =
+                        getCustomAppTransition(animAttr, container);
+                if (customAppTransition != null) {
+                    a = loadCustomActivityAnimation(customAppTransition, enter, container);
+                } else {
                     if (canCustomizeAppTransition) {
                         a = loadAnimationAttr(lp, animAttr, transit);
                     } else {
@@ -911,7 +914,7 @@ public class AppTransition implements Dump {
         return a;
     }
 
-    Animation loadCustomActivityAnimation(int animAttr, boolean enter, WindowContainer container) {
+    CustomAppTransition getCustomAppTransition(int animAttr, WindowContainer container) {
         ActivityRecord customAnimationSource = container.asActivityRecord();
         if (customAnimationSource == null) {
             return null;
@@ -927,31 +930,28 @@ public class AppTransition implements Dump {
                 return null;
             }
         }
-        final CustomAppTransition custom;
         switch (animAttr) {
             case WindowAnimation_activityOpenEnterAnimation:
             case WindowAnimation_activityOpenExitAnimation:
-                custom = customAnimationSource.getCustomAnimation(true /* open */);
-                break;
+                return customAnimationSource.getCustomAnimation(true /* open */);
             case WindowAnimation_activityCloseEnterAnimation:
             case WindowAnimation_activityCloseExitAnimation:
-                custom = customAnimationSource.getCustomAnimation(false /* open */);
-                break;
-            default:
-                return null;
-        }
-        if (custom != null) {
-            final Animation a = mTransitionAnimation.loadAppTransitionAnimation(
-                    customAnimationSource.packageName, enter
-                            ? custom.mEnterAnim : custom.mExitAnim);
-            if (a != null && custom.mBackgroundColor != 0) {
-                a.setBackdropColor(custom.mBackgroundColor);
-                a.setShowBackdrop(true);
-            }
-            return a;
+                return customAnimationSource.getCustomAnimation(false /* open */);
         }
         return null;
     }
+    private Animation loadCustomActivityAnimation(@NonNull CustomAppTransition custom,
+            boolean enter, WindowContainer container) {
+        final ActivityRecord customAnimationSource = container.asActivityRecord();
+        final Animation a = mTransitionAnimation.loadAppTransitionAnimation(
+                customAnimationSource.packageName, enter
+                        ? custom.mEnterAnim : custom.mExitAnim);
+        if (a != null && custom.mBackgroundColor != 0) {
+            a.setBackdropColor(custom.mBackgroundColor);
+            a.setShowBackdrop(true);
+        }
+        return a;
+    }
 
     int getAppRootTaskClipMode() {
         return mNextAppTransitionRequests.contains(TRANSIT_RELAUNCH)