Fix deletion of custom block classes

packages/editor/src/hooks/custom-class-name.js

@@ -1,7 +1,7 @@
 /**
  * External dependencies
  */
-import { assign, difference, compact } from 'lodash';
+import { assign, difference, omit } from 'lodash';
 import classnames from 'classnames';
 
 /**
@@ -134,20 +134,17 @@ export function getHTMLRootElementClasses( innerHTML ) {
 export function addParsedDifference( blockAttributes, blockType, innerHTML ) {
 	if ( hasBlockSupport( blockType, 'customClassName', true ) ) {
 		// To determine difference, serialize block given the known set of
-		// attributes. If there are classes which are mismatched with the
-		// incoming HTML of the block, add to filtered result.
-		const serialized = getSaveContent( blockType, blockAttributes );
-		const classes = getHTMLRootElementClasses( serialized );
-		const parsedClasses = getHTMLRootElementClasses( innerHTML );
-		const customClasses = difference( parsedClasses, classes );
-
-		const filteredClassName = compact( [
-			blockAttributes.className,
-			...customClasses,
-		] ).join( ' ' );
-
-		if ( filteredClassName ) {
-			blockAttributes.className = filteredClassName;
+		// attributes, with the exception of `className`. This will determine
+		// the default set of classes. From there, any difference in innerHTML
+		// can be considered as custom classes.
+		const attributesSansClassName = omit( blockAttributes, [ 'className' ] );
+		const serialized = getSaveContent( blockType, attributesSansClassName );
+		const defaultClasses = getHTMLRootElementClasses( serialized );
+		const actualClasses = getHTMLRootElementClasses( innerHTML );
+		const customClasses = difference( actualClasses, defaultClasses );
+
+		if ( customClasses.length ) {
+			blockAttributes.className = customClasses.join( ' ' );
 		} else {
 			delete blockAttributes.className;
 		}

packages/editor/src/hooks/test/custom-class-name.js

@@ -10,7 +10,7 @@ import { getHTMLRootElementClasses } from '../custom-class-name';
 
 describe( 'custom className', () => {
 	const blockSettings = {
-		save: () => <div />,
+		save: () => <div className="default" />,
 		category: 'common',
 		title: 'block title',
 	};
@@ -95,7 +95,7 @@ describe( 'custom className', () => {
 			const attributes = addParsedDifference(
 				{ className: 'foo' },
 				blockSettings,
-				'<div class="foo bar baz"></div>'
+				'<div class="default foo bar baz"></div>'
 			);
 
 			expect( attributes.className ).toBe( 'foo bar baz' );
@@ -110,5 +110,45 @@ describe( 'custom className', () => {
 
 			expect( attributes.className ).toBeUndefined();
 		} );
+
+		it( 'should add a custom class name to an element without a class', () => {
+			const attributes = addParsedDifference(
+				{},
+				blockSettings,
+				'<div class="default foo"></div>'
+			);
+
+			expect( attributes.className ).toBe( 'foo' );
+		} );
+
+		it( 'should remove the custom class and retain default class', () => {
+			const attributes = addParsedDifference(
+				{ className: 'custom1 custom2' },
+				blockSettings,
+				'<div class="default custom1"></div>'
+			);
+
+			expect( attributes.className ).toBe( 'custom1' );
+		} );
+
+		it( 'should remove the custom class from an element originally without a class', () => {
+			const attributes = addParsedDifference(
+				{ className: 'foo' },
+				blockSettings,
+				'<div></div>'
+			);
+
+			expect( attributes.className ).toBeUndefined();
+		} );
+
+		it( 'should remove the custom classes and retain default and other custom class', () => {
+			const attributes = addParsedDifference(
+				{ className: 'custom1 custom2 custom3' },
+				blockSettings,
+				'<div class="default custom1 custom3"></div>'
+			);
+
+			expect( attributes.className ).toBe( 'custom1 custom3' );
+		} );
 	} );
 } );