store a reference to all added controls, and remove them when the map…

… is removed simpler loop + new test update test
mapbox · Jul 30, 2018 · 4d452cd · 4d452cd
1 parent 8f8f382
commit 4d452cd
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/src/ui/map.js b/src/ui/map.js
@@ -259,6 +259,7 @@ class Map extends Camera {
     _crossFadingFactor: number;
     _collectResourceTiming: boolean;
     _renderTaskQueue: TaskQueue;
+    _controls: Array<IControl>;
 
     /**
      * The map's {@link ScrollZoomHandler}, which implements zooming in and out with a scroll wheel or trackpad.
@@ -319,6 +320,7 @@ class Map extends Camera {
         this._crossFadingFactor = 1;
         this._collectResourceTiming = options.collectResourceTiming;
         this._renderTaskQueue = new TaskQueue();
+        this._controls = [];
 
         const transformRequestFn = options.transformRequest;
         this._transformRequest = transformRequestFn ?  (url, type) => transformRequestFn(url, type) || ({ url }) : (url) => ({ url });
@@ -411,7 +413,13 @@ class Map extends Camera {
         if (position === undefined) {
             position = 'top-right';
         }
+        if (!control || !control.onAdd) {
+            return this.fire(new ErrorEvent(new Error(
+                'Invalid argument to map.addControl(). Argument must be a control with onAdd and onRemove methods.')));
+        }
         const controlElement = control.onAdd(this);
+        this._controls.push(control);
+
         const positionContainer = this._controlPositions[position];
         if (position.indexOf('bottom') !== -1) {
             positionContainer.insertBefore(controlElement, positionContainer.firstChild);
@@ -428,6 +436,12 @@ class Map extends Camera {
      * @returns {Map} `this`
      */
     removeControl(control: IControl) {
+        if (!control || !control.onRemove) {
+            return this.fire(new ErrorEvent(new Error(
+                'Invalid argument to map.removeControl(). Argument must be a control with onAdd and onRemove methods.')));
+        }
+        const ci = this._controls.indexOf(control);
+        if (ci > -1) this._controls.splice(ci, 1);
         control.onRemove(this);
         return this;
     }
@@ -1723,6 +1737,10 @@ class Map extends Camera {
             window.removeEventListener('resize', this._onWindowResize, false);
             window.removeEventListener('online', this._onWindowOnline, false);
         }
+
+        for (const control of this._controls) control.onRemove(this);
+        this._controls = [];
+
         const extension = this.painter.context.gl.getExtension('WEBGL_lose_context');
         if (extension) extension.loseContext();
         removeNode(this._canvasContainer);

diff --git a/test/unit/ui/map.test.js b/test/unit/ui/map.test.js
@@ -790,18 +790,46 @@ test('Map', (t) => {
         t.end();
     });
 
+    t.test('#remove calls onRemove on added controls', (t) => {
+        const map = createMap(t);
+        const control = {
+            onRemove: t.spy(),
+            onAdd: function (_) {
+                return window.document.createElement('div');
+            }
+        };
+        map.addControl(control);
+        map.remove();
+        t.ok(control.onRemove.calledOnce);
+        t.end();
+    });
+
     t.test('#addControl', (t) => {
         const map = createMap(t);
         const control = {
             onAdd: function(_) {
                 t.equal(map, _, 'addTo() called with map');
-                t.end();
                 return window.document.createElement('div');
             }
         };
         map.addControl(control);
+        t.equal(map._controls[1], control, "saves reference to added controls");
+        t.end();
+    });
+
+    t.test('#removeControl errors on invalid arguments', (t) => {
+        const map = createMap(t);
+        const control = {};
+        const stub = t.stub(console, 'error');
+
+        map.addControl(control);
+        map.removeControl(control);
+        t.ok(stub.calledTwice);
+        t.end();
+
     });
 
+
     t.test('#removeControl', (t) => {
         const map = createMap(t);
         const control = {
@@ -810,11 +838,13 @@ test('Map', (t) => {
             },
             onRemove: function(_) {
                 t.equal(map, _, 'onRemove() called with map');
-                t.end();
             }
         };
         map.addControl(control);
         map.removeControl(control);
+        t.equal(map._controls.length, 1, "removes removed controls from map's control array");
+        t.end();
+
     });
 
     t.test('#project', (t) => {