ObservableMap: optimize unsubscribe and add unit tests

packages/components/src/slot-fill/bubbles-virtually/observable-map.ts

@@ -10,6 +10,10 @@ export type ObservableMap< K, V > = {
 	subscribe( name: K, listener: () => void ): () => void;
 };
 
+/**
+ * A key/value map where the individual entries are observable by subscribing to them
+ * with the `subscribe` methods.
+ */
 export function observableMap< K, V >(): ObservableMap< K, V > {
 	const map = new Map< K, V >();
 	const listeners = new Map< K, Set< () => void > >();
@@ -24,20 +28,6 @@ export function observableMap< K, V >(): ObservableMap< K, V > {
 		}
 	}
 
-	function unsubscribe( name: K, listener: () => void ) {
-		return () => {
-			const list = listeners.get( name );
-			if ( ! list ) {
-				return;
-			}
-
-			list.delete( listener );
-			if ( list.size === 0 ) {
-				listeners.delete( name );
-			}
-		};
-	}
-
 	return {
 		get( name ) {
 			return map.get( name );
@@ -58,11 +48,22 @@ export function observableMap< K, V >(): ObservableMap< K, V > {
 			}
 			list.add( listener );
 
-			return unsubscribe( name, listener );
+			return () => {
+				list.delete( listener );
+				if ( list.size === 0 ) {
+					listeners.delete( name );
+				}
+			};
 		},
 	};
 }
 
+/**
+ * React hook that lets you observe an individual entry in an `ObservableMap`.
+ *
+ * @param map  The `ObservableMap` to observe.
+ * @param name The map key to observe.
+ */
 export function useObservableValue< K, V >(
 	map: ObservableMap< K, V >,
 	name: K

packages/components/src/slot-fill/test/observable-map.js

@@ -0,0 +1,83 @@
+/**
+ * External dependencies
+ */
+import { render, screen, act } from '@testing-library/react';
+
+/**
+ * Internal dependencies
+ */
+import {
+	observableMap,
+	useObservableValue,
+} from '../bubbles-virtually/observable-map';
+
+describe( 'ObservableMap', () => {
+	test( 'should observe individual values', () => {
+		const map = observableMap();
+
+		const listenerA = jest.fn();
+		const listenerB = jest.fn();
+
+		const unsubA = map.subscribe( 'a', listenerA );
+		const unsubB = map.subscribe( 'b', listenerB );
+
+		// check that setting `a` doesn't notify the `b` listener
+		map.set( 'a', 1 );
+		expect( listenerA ).toHaveBeenCalledTimes( 1 );
+		expect( listenerB ).toHaveBeenCalledTimes( 0 );
+
+		// check that setting `b` doesn't notify the `a` listener
+		map.set( 'b', 2 );
+		expect( listenerA ).toHaveBeenCalledTimes( 1 );
+		expect( listenerB ).toHaveBeenCalledTimes( 1 );
+
+		// check that `delete` triggers notifications, too
+		map.delete( 'a' );
+		expect( listenerA ).toHaveBeenCalledTimes( 2 );
+		expect( listenerB ).toHaveBeenCalledTimes( 1 );
+
+		// check that the subscription survived the `delete`
+		map.set( 'a', 2 );
+		expect( listenerA ).toHaveBeenCalledTimes( 3 );
+		expect( listenerB ).toHaveBeenCalledTimes( 1 );
+
+		// check that unsubscription really works
+		unsubA();
+		unsubB();
+		map.set( 'a', 3 );
+		expect( listenerA ).toHaveBeenCalledTimes( 3 );
+		expect( listenerB ).toHaveBeenCalledTimes( 1 );
+	} );
+} );
+
+describe( 'useObservableValue', () => {
+	test( 'reacts only to the specified key', () => {
+		const map = observableMap();
+		map.set( 'a', 1 );
+
+		const MapUI = jest.fn( () => {
+			const value = useObservableValue( map, 'a' );
+			return <div>value is { value }</div>;
+		} );
+
+		render( <MapUI /> );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 1'
+		);
+		expect( MapUI ).toHaveBeenCalledTimes( 1 );
+
+		act( () => {
+			map.set( 'a', 2 );
+		} );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 2'
+		);
+		expect( MapUI ).toHaveBeenCalledTimes( 2 );
+
+		// check that setting unobserved map key doesn't trigger a render at all
+		act( () => {
+			map.set( 'b', 1 );
+		} );
+		expect( MapUI ).toHaveBeenCalledTimes( 2 );
+	} );
+} );