fix: datasource context provider will now cleanup when being unmounted

__mocks__/azure-maps-control.js

@@ -1,9 +1,18 @@
 class DataSource {
+  id
+  options
+
+  constructor(id, options) {
+    this.id = id
+    this.options = options
+  }
+
   add = jest.fn()
   clear = jest.fn()
   remove = jest.fn()
   importDataFromUrl = jest.fn()
-  setOptions = jest.fn()
+  setOptions = jest.fn((options) => (this.options = options))
+  getId = () => this.id
 }
 
 module.exports = {
@@ -12,9 +21,14 @@ module.exports = {
       add: jest.fn()
     },
     events: {
-      add: jest.fn((eventName, callback = () => {}) => {
-        callback()
-      })
+      add: jest.fn((_eventName, _targetOrCallback, callback = () => {}) => {
+        if (typeof _targetOrCallback === 'function') {
+          _targetOrCallback()
+        } else {
+          callback()
+        }
+      }),
+      remove: jest.fn((eventName) => {})
     },
     imageSprite: {
       add: jest.fn(),
@@ -26,7 +40,9 @@ module.exports = {
     },
     layers: {
       add: jest.fn(),
-      remove: jest.fn()
+      remove: jest.fn(),
+      getLayers: jest.fn(() => []),
+      getLayerById: jest.fn()
     },
     popups: {
       getPopups: jest.fn(() => []),
@@ -123,7 +139,7 @@ module.exports = {
     VectorTileSource: jest.fn((id, options) => ({
       getId: jest.fn(() => id),
       getOptions: jest.fn(() => options)
-    })) 
+    }))
   },
   Shape: jest.fn(() => ({
     setCoordinates: jest.fn(),
@@ -132,7 +148,7 @@ module.exports = {
   data: {
     Position: jest.fn((...args) => args),
     BoundingBox: jest.fn((...args) => args),
-    Point: jest.fn(coords => ({ coords, type: 'Point' })),
+    Point: jest.fn((coords) => ({ coords, type: 'Point' })),
     MultiPoint: jest.fn((coords, bbox) => ({ coords, bbox, type: 'MultiPoint' })),
     LineString: jest.fn((coords, bbox) => ({ coords, bbox, type: 'LineString' })),
     MultiLineString: jest.fn((multipleCoordinates, bbox) => ({

src/components/helpers/mapHelper.ts

@@ -1,6 +1,9 @@
 import atlas from 'azure-maps-control'
+import { DataSourceType, MapType } from '../../types'
 
-export const generateLinesFromArrayOfPosition = (coordinates: atlas.data.Position[]): atlas.data.LineString => {
+export const generateLinesFromArrayOfPosition = (
+  coordinates: atlas.data.Position[]
+): atlas.data.LineString => {
   const line = new atlas.data.LineString(coordinates)
   return line
 }
@@ -9,3 +12,14 @@ export const generatePixelHeading = (origin: atlas.Pixel, destination: atlas.Pix
   const heading = atlas.Pixel.getHeading(origin, destination)
   return heading
 }
+
+export const getLayersDependingOnDatasource = (mref: MapType, dst: DataSourceType) => {
+  return mref.layers.getLayers().filter((l) => {
+    if ((l as atlas.layer.SymbolLayer).getSource) {
+      const sourceLayer = (l as atlas.layer.SymbolLayer).getSource()
+      const dsId = typeof sourceLayer === 'string' ? sourceLayer : sourceLayer.getId()
+      return dsId === dst.getId()
+    }
+    return false
+  })
+}

src/contexts/AzureMapDataSourceContext.test.tsx

@@ -1,6 +1,6 @@
 import { useContext } from 'react'
 import { renderHook } from '@testing-library/react'
-import atlas, { Map } from 'azure-maps-control'
+import atlas, { layer, Map } from 'azure-maps-control'
 import React from 'react'
 import { AzureMapsContext } from '../contexts/AzureMapContext'
 import {
@@ -93,12 +93,49 @@ describe('AzureMapDataSourceProvider tests', () => {
   })
 
   it('should call setOptions and clear method if options was changed', () => {
-    const { result, rerender } = renderHook(() => useContextConsumer(), {
+    const { result } = renderHook(() => useContextConsumer(), {
       wrapper: wrapWithDataSourceContext({ id: 'id', options: { option: 'option' } })
     })
     expect(result.current.dataSourceRef).toBeInstanceOf(atlas.source.DataSource)
     expect(
       (result.current.dataSourceRef as atlas.source.DataSource).setOptions
     ).toHaveBeenLastCalledWith({ option: 'option' })
   })
+
+  it('should remove data source from the map ref on unmount', () => {
+    mapRef.events.remove = jest.fn()
+    const events = { render: () => {} }
+    const { unmount, result } = renderHook(() => useContextConsumer(), {
+      wrapper: wrapWithDataSourceContext({ id: 'id', options: { option: 'option' }, events })
+    })
+
+    unmount()
+
+    expect(mapRef.sources.remove).toHaveBeenCalledWith(result.current.dataSourceRef)
+    expect(mapRef.events.remove).toHaveBeenCalledWith(
+      'render',
+      result.current.dataSourceRef,
+      events.render
+    )
+  })
+
+  it('should remove all layers that are using the same datasource from the map ref on unmount', () => {
+    const dsToBeRemoved = new atlas.source.DataSource('ds_to_be_removed')
+    const dsToKeep = new atlas.source.DataSource('ds_to_keep')
+
+    const symbolLayer = new layer.SymbolLayer(dsToBeRemoved, 'layer_to_be_removed')
+    const bubbleLayer = new layer.BubbleLayer(dsToKeep, 'layer_to_keep')
+
+    symbolLayer.getSource = jest.fn(() => dsToBeRemoved)
+
+    mapRef.layers.getLayers = jest.fn(() => [symbolLayer, bubbleLayer])
+
+    const { unmount } = renderHook(() => useContextConsumer(), {
+      wrapper: wrapWithDataSourceContext({ id: dsToBeRemoved.getId() })
+    })
+
+    unmount()
+    expect(mapRef.layers.remove).toHaveBeenCalledTimes(1)
+    expect(mapRef.layers.remove).toHaveBeenNthCalledWith(1, 'layer_to_be_removed')
+  })
 })

src/contexts/AzureMapDataSourceContext.tsx

@@ -9,6 +9,7 @@ import {
 import atlas from 'azure-maps-control'
 import { AzureMapsContext } from './AzureMapContext'
 import { useCheckRef } from '../hooks/useCheckRef'
+import { getLayersDependingOnDatasource } from '../components/helpers/mapHelper'
 
 const AzureMapDataSourceContext = createContext<IAzureMapDataSourceProps>({
   dataSourceRef: null
@@ -31,7 +32,9 @@ const AzureMapDataSourceStatefulProvider = ({
   dataFromUrl,
   collection
 }: IAzureDataSourceStatefulProviderProps) => {
-  const [dataSourceRef] = useState<atlas.source.DataSource>(new atlas.source.DataSource(id, options))
+  const [dataSourceRef] = useState<atlas.source.DataSource>(
+    new atlas.source.DataSource(id, options)
+  )
   const { mapRef } = useContext<IAzureMapsContextProps>(AzureMapsContext)
   useCheckRef<MapType, DataSourceType>(mapRef, dataSourceRef, (mref, dref) => {
     for (const eventType in events || {}) {
@@ -46,6 +49,16 @@ const AzureMapDataSourceStatefulProvider = ({
         dref.add(collection)
       }
     }
+    return () => {
+      for (const eventType in events || {}) {
+        mref.events.remove(eventType as any, dref, events[eventType])
+      }
+
+      getLayersDependingOnDatasource(mref, dref).forEach((l) => {
+        mref.layers.remove(l.getId() ? l.getId() : l)
+      })
+      mref.sources.remove(dref)
+    }
   })
 
   useEffect(() => {

src/contexts/AzureMapVectorTileSourceProvider.test.tsx

@@ -1,6 +1,6 @@
 import { renderHook } from '@testing-library/react'
 import React, { useContext } from 'react'
-import { Map } from 'azure-maps-control'
+import atlas, { layer, Map } from 'azure-maps-control'
 import { IAzureVectorTileSourceStatefulProviderProps } from '../types'
 import { AzureMapsContext } from './AzureMapContext'
 import { AzureMapVectorTileSourceProvider } from './AzureMapVectorTileSourceProvider'
@@ -71,4 +71,41 @@ describe('AzureMapVectorTileSourceProvider tests', () => {
       expect.any(Function)
     )
   })
+
+  it('should remove data source from the map ref on unmount', () => {
+    mapRef.events.remove = jest.fn()
+    const events = { sourceadded: () => {} }
+    const { unmount, result } = renderHook(() => useContextConsumer(), {
+      wrapper: wrapWithVectorTileSourceContext({ id: 'id', options: { option: 'option' }, events })
+    })
+
+    unmount()
+
+    expect(mapRef.sources.remove).toHaveBeenCalledWith(result.current.dataSourceRef)
+    expect(mapRef.events.remove).toHaveBeenCalledWith(
+      'sourceadded',
+      result.current.dataSourceRef,
+      events.sourceadded
+    )
+  })
+
+  it('should remove all layers that are using the same datasource from the map ref on unmount', () => {
+    const dsToBeRemoved = new atlas.source.DataSource('ds_to_be_removed')
+    const dsToKeep = new atlas.source.DataSource('ds_to_keep')
+
+    const symbolLayer = new layer.SymbolLayer(dsToBeRemoved, 'layer_to_be_removed')
+    const bubbleLayer = new layer.BubbleLayer(dsToKeep, 'layer_to_keep')
+
+    symbolLayer.getSource = jest.fn(() => dsToBeRemoved)
+
+    mapRef.layers.getLayers = jest.fn(() => [symbolLayer, bubbleLayer])
+
+    const { unmount } = renderHook(() => useContextConsumer(), {
+      wrapper: wrapWithVectorTileSourceContext({ id: dsToBeRemoved.getId() })
+    })
+
+    unmount()
+    expect(mapRef.layers.remove).toHaveBeenCalledTimes(1)
+    expect(mapRef.layers.remove).toHaveBeenNthCalledWith(1, 'layer_to_be_removed')
+  })
 })

src/contexts/AzureMapVectorTileSourceProvider.tsx

@@ -1,9 +1,16 @@
 import atlas from 'azure-maps-control'
 import React, { useContext, useState } from 'react'
 import { useCheckRef } from '../hooks/useCheckRef'
-import { DataSourceType, IAzureMapsContextProps, IAzureMapSourceEventType, IAzureVectorTileSourceStatefulProviderProps, MapType } from '../types'
+import {
+  DataSourceType,
+  IAzureMapsContextProps,
+  IAzureMapSourceEventType,
+  IAzureVectorTileSourceStatefulProviderProps,
+  MapType
+} from '../types'
 import { AzureMapDataSourceRawProvider as Provider } from './AzureMapDataSourceContext'
 import { AzureMapsContext } from './AzureMapContext'
+import { getLayersDependingOnDatasource } from '../components/helpers/mapHelper'
 
 /**
  * @param id datasource identifier
@@ -15,29 +22,49 @@ const AzureMapVectorTileSourceStatefulProvider = ({
   id,
   children,
   options,
-  events = {},
+  events = {}
 }: IAzureVectorTileSourceStatefulProviderProps) => {
-  const [dataSourceRef] = useState<atlas.source.VectorTileSource>(new atlas.source.VectorTileSource(id, options))
+  const [dataSourceRef] = useState<atlas.source.VectorTileSource>(
+    new atlas.source.VectorTileSource(id, options)
+  )
   const { mapRef } = useContext<IAzureMapsContextProps>(AzureMapsContext)
   useCheckRef<MapType, DataSourceType>(mapRef, dataSourceRef, (mref, dref) => {
     for (const eventType in events) {
-      const handler = events[eventType as IAzureMapSourceEventType] as (e: atlas.source.Source) => void | undefined
-      if(handler) {
+      const handler = events[eventType as IAzureMapSourceEventType] as (
+        e: atlas.source.Source
+      ) => void | undefined
+      if (handler) {
         mref.events.add(eventType as IAzureMapSourceEventType, dref, handler)
       }
     }
     mref.sources.add(dref)
+
+    return () => {
+      for (const eventType in events || {}) {
+        const handler = events[eventType as IAzureMapSourceEventType] as (
+          e: atlas.source.Source
+        ) => void | undefined
+        if (handler) {
+          mref.events.remove(eventType as IAzureMapSourceEventType, dref, handler)
+        }
+      }
+
+      getLayersDependingOnDatasource(mref, dref).forEach((l) => {
+        mref.layers.remove(l.getId() ? l.getId() : l)
+      })
+      mref.sources.remove(dref)
+    }
   })
 
   return (
     <Provider
       value={{
-        dataSourceRef,
+        dataSourceRef
       }}
     >
       {mapRef && children}
     </Provider>
   )
 }
 
-export { AzureMapVectorTileSourceStatefulProvider as AzureMapVectorTileSourceProvider }
+export { AzureMapVectorTileSourceStatefulProvider as AzureMapVectorTileSourceProvider }

src/hooks/useAzureMapLayer.test.tsx

@@ -1,12 +1,12 @@
-import atlas, { source, layer } from 'azure-maps-control'
+import { source, layer } from 'azure-maps-control'
 import { ReactNode } from 'react'
 import { renderHook } from '@testing-library/react'
 import { useAzureMapLayer } from './useAzureMapLayer'
 import { Map } from 'azure-maps-control'
 import React from 'react'
 import { AzureMapsContext } from '../contexts/AzureMapContext'
 import { AzureMapDataSourceContext } from '../contexts/AzureMapDataSourceContext'
-import { IAzureLayerStatefulProviderProps, LayerType } from '../types'
+import { IAzureLayerStatefulProviderProps } from '../types'
 
 const mapContextProps = {
   mapRef: null,
@@ -16,6 +16,7 @@ const mapContextProps = {
   setMapRef: jest.fn()
 }
 const mapRef = new Map('fake', {})
+mapRef.layers.getLayerById = jest.fn().mockImplementation(() => null)
 
 const wrapWithAzureMapContext = ({ children }: { children?: ReactNode | null }) => {
   const datasourceRef = {} as source.DataSource
@@ -109,6 +110,8 @@ describe('useAzureMapLayer tests', () => {
   })
 
   it('shouldRemove layer from map on unmoun', () => {
+    const symbolLayer = {} as layer.SymbolLayer
+    mapRef.layers.getLayerById = jest.fn().mockImplementation(() => symbolLayer)
     mapRef.layers.remove = jest.fn()
     const { unmount } = renderHook(
       () =>

src/hooks/useAzureMapLayer.tsx

@@ -70,7 +70,9 @@ export const useAzureMapLayer = ({
     mref.layers.add(lref)
     return () => {
       try {
-        mref.layers.remove(lref.getId() ? lref.getId() : lref)
+        if (mref.layers.getLayerById(lref.getId())) {
+          mref.layers.remove(lref.getId() ? lref.getId() : lref)
+        }
       } catch (e) {
         console.error('Error on remove layer', e)
       }