From 7a8ed11e5f9d7af3f210836d10ea0672b42aa998 Mon Sep 17 00:00:00 2001
From: Yuya Tanaka <yuya.presto@gmail.com>
Date: Thu, 26 Dec 2019 04:18:57 +0900
Subject: [PATCH] [Badge] Retain text, color and variant while hide transition

---
 packages/material-ui/src/Badge/Badge.js      | 26 ++++++++++++++------
 packages/material-ui/src/Badge/Badge.test.js | 15 +++++++++++
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/packages/material-ui/src/Badge/Badge.js b/packages/material-ui/src/Badge/Badge.js
index cf1fcef48e5a8e..bc47f0cad49bd6 100644
--- a/packages/material-ui/src/Badge/Badge.js
+++ b/packages/material-ui/src/Badge/Badge.js
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useRef } from 'react';
 import PropTypes from 'prop-types';
 import clsx from 'clsx';
 import withStyles from '../styles/withStyles';
@@ -161,13 +161,13 @@ const Badge = React.forwardRef(function Badge(props, ref) {
     children,
     classes,
     className,
-    color = 'default',
+    color: colorProp = 'default',
     component: ComponentProp = 'span',
     invisible: invisibleProp,
     max = 99,
     overlap = 'rectangle',
     showZero = false,
-    variant = 'standard',
+    variant: variantProp = 'standard',
     ...other
   } = props;
 
@@ -175,17 +175,29 @@ const Badge = React.forwardRef(function Badge(props, ref) {
 
   if (
     invisibleProp == null &&
-    ((badgeContent === 0 && !showZero) || (badgeContent == null && variant !== 'dot'))
+    ((badgeContent === 0 && !showZero) || (badgeContent == null && variantProp !== 'dot'))
   ) {
     invisible = true;
   }
 
-  let displayValue = '';
+  let nextDisplayValue = '';
 
-  if (variant !== 'dot') {
-    displayValue = badgeContent > max ? `${max}+` : badgeContent;
+  if (variantProp !== 'dot') {
+    nextDisplayValue = badgeContent > max ? `${max}+` : badgeContent;
   }
 
+  // Retain the appearance of badge while invisible, to keep same appearance until disappearing.
+  // These should not be an attribute with transitions, or it will cause an undesirable animation at next appearing.
+  const nextNotTransitionedAttrs = {
+    displayValue: nextDisplayValue,
+    color: colorProp,
+    variant: variantProp,
+  };
+  const lastNotTransitionedPropsRef = useRef(nextNotTransitionedAttrs);
+  const { displayValue, color, variant } = invisible
+    ? lastNotTransitionedPropsRef.current
+    : nextNotTransitionedAttrs;
+
   return (
     <ComponentProp className={clsx(classes.root, className)} ref={ref} {...other}>
       {children}
diff --git a/packages/material-ui/src/Badge/Badge.test.js b/packages/material-ui/src/Badge/Badge.test.js
index 80bcc58c31c770..01e8250dc05964 100644
--- a/packages/material-ui/src/Badge/Badge.test.js
+++ b/packages/material-ui/src/Badge/Badge.test.js
@@ -3,6 +3,7 @@ import { assert } from 'chai';
 import { createMount, getClasses } from '@material-ui/core/test-utils';
 import describeConformance from '../test-utils/describeConformance';
 import Badge from './Badge';
+import { act } from 'react-dom/test-utils';
 
 function findBadge(wrapper) {
   return wrapper.find('span').at(1);
@@ -57,6 +58,20 @@ describe('<Badge />', () => {
     assert.strictEqual(wrapper.contains(defaultProps.children), true);
   });
 
+  it('retains text, color and variant while invisible for disappearing transition', () => {
+    const wrapper = mount(<Badge {...defaultProps} color="primary" variant="dot" />);
+    act(() => {
+      wrapper.setProps({ badgeContent: 0, color: 'secondary', variant: 'standard' });
+    });
+    assert.strictEqual(findBadge(wrapper).text(), '');
+    assert.strictEqual(findBadge(wrapper).hasClass(classes.colorPrimary), true);
+    assert.strictEqual(findBadge(wrapper).hasClass(classes.dot), true);
+    act(() => {
+      wrapper.setProps({ showZero: true });
+    });
+    assert.strictEqual(findBadge(wrapper).text(), '0');
+  });
+
   describe('prop: color', () => {
     it('should have the colorPrimary class when color="primary"', () => {
       const wrapper = mount(<Badge {...defaultProps} color="primary" />);