ObservableMap: optimize unsubscribe and add unit tests

Author: jsnajdr

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,18 +28,14 @@ export function observableMap< K, V >(): ObservableMap< K, V > {
 		}
 	}
 
-	function unsubscribe( name: K, listener: () => void ) {
-		return () => {
-			const list = listeners.get( name );
-			if ( ! list ) {
-				return;
-			}
+	function getListeners( name: K ) {
+		let list = listeners.get( name );
+		if ( ! list ) {
+			list = new Set();
+			listeners.set( name, list );
+		}
 
-			list.delete( listener );
-			if ( list.size === 0 ) {
-				listeners.delete( name );
-			}
-		};
+		return list;
 	}
 
 	return {
@@ -51,18 +51,25 @@ export function observableMap< K, V >(): ObservableMap< K, V > {
 			callListeners( name );
 		},
 		subscribe( name, listener ) {
-			let list = listeners.get( name );
-			if ( ! list ) {
-				list = new Set();
-				listeners.set( name, list );
-			}
+			const list = getListeners( name );
 			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,87 @@
+/**
+ * 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();
+		let notifsFromA = 0;
+		let notifsFromB = 0;
+
+		const unsubA = map.subscribe( 'a', () => {
+			notifsFromA++;
+		} );
+		const unsubB = map.subscribe( 'b', () => {
+			notifsFromB++;
+		} );
+
+		// check that setting `a` doesn't notify the `b` listener
+		map.set( 'a', 1 );
+		expect( notifsFromA ).toBe( 1 );
+		expect( notifsFromB ).toBe( 0 );
+
+		// check that setting `b` doesn't notify the `a` listener
+		map.set( 'b', 2 );
+		expect( notifsFromA ).toBe( 1 );
+		expect( notifsFromB ).toBe( 1 );
+
+		// check that `delete` triggers notifications, too
+		map.delete( 'a' );
+		expect( notifsFromA ).toBe( 2 );
+		expect( notifsFromB ).toBe( 1 );
+
+		// check that the subscription survived the `delete`
+		map.set( 'a', 2 );
+		expect( notifsFromA ).toBe( 3 );
+		expect( notifsFromB ).toBe( 1 );
+
+		// check that unsubscription really works.
+		unsubA();
+		unsubB();
+		map.set( 'a', 3 );
+		expect( notifsFromA ).toBe( 3 );
+		expect( notifsFromB ).toBe( 1 );
+	} );
+} );
+
+describe( 'useObservableValue', () => {
+	test( 'reacts only to the specified key', () => {
+		const map = observableMap();
+		map.set( 'a', 1 );
+
+		let renders = 0;
+		function MapUI() {
+			const value = useObservableValue( map, 'a' );
+			renders++;
+			return <div>value is { value }</div>;
+		}
+
+		render( <MapUI /> );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 1'
+		);
+		expect( renders ).toBe( 1 );
+
+		act( () => {
+			map.set( 'a', 2 );
+		} );
+		expect( screen.getByText( /^value is/ ) ).toHaveTextContent(
+			'value is 2'
+		);
+		expect( renders ).toBe( 2 );
+
+		act( () => {
+			map.set( 'b', 1 );
+		} );
+		expect( renders ).toBe( 2 );
+	} );
+} );