diff --git a/dftimewolf/lib/containers/manager.py b/dftimewolf/lib/containers/manager.py index 29ce80f7d..6ccfcbb1f 100644 --- a/dftimewolf/lib/containers/manager.py +++ b/dftimewolf/lib/containers/manager.py @@ -129,7 +129,8 @@ def StoreContainer(self, module.storage[container.CONTAINER_TYPE] = [] # If the container to add exists already in the state, don't add it again - if container in [c for c, _ in module.storage[container.CONTAINER_TYPE]]: + if any(c is container for c, _ in module.storage[container.CONTAINER_TYPE]): + self._logger.debug('Duplicate container not added to storage: %s', str(container)) continue module.storage[container.CONTAINER_TYPE].append((container, source_module)) diff --git a/tests/lib/containers/manager.py b/tests/lib/containers/manager.py index a212d4447..6182bbb69 100644 --- a/tests/lib/containers/manager.py +++ b/tests/lib/containers/manager.py @@ -4,8 +4,6 @@ import unittest from unittest import mock -import pandas as pd - from dftimewolf.lib.containers import containers from dftimewolf.lib.containers import interface from dftimewolf.lib.containers import manager @@ -548,67 +546,20 @@ def test_StoreDuplicateContainers(self): """Tests that attempts to store duplicate containers are disregarded.""" self._container_manager.ParseRecipe(_TEST_RECIPE) - self._container_manager.StoreContainer( - source_module='Preflight1', container=_TestContainer1('param1')) - self._container_manager.StoreContainer( - source_module='Preflight1', container=_TestContainer1('param1')) - self._container_manager.StoreContainer( - source_module='Preflight1', container=_TestContainer2('param1')) - self._container_manager.StoreContainer( - source_module='Preflight1', container=_TestContainer2('param1')) - self._container_manager.StoreContainer( - source_module='Preflight1', container=_TestContainer2('param2')) + container_1 = _TestContainer1('param1') + container_2 = _TestContainer2('param1') + container_3 = _TestContainer2('param2') - # Metadata is not considered in duplicate comparison - c5 = _TestContainer3('param1') - c6 = _TestContainer3('param1') - c5.SetMetadata('key', 'foo') - c6.SetMetadata('key', 'bar') self._container_manager.StoreContainer( - source_module='Preflight1', container=c5) + source_module='Preflight1', container=container_1) self._container_manager.StoreContainer( - source_module='Preflight1', container=c6) - - # Dataframe members of containers have special handling; check that too - df1 = pd.DataFrame(columns=['a', 'b'], data=[[1, 2], [3, 4]]) - df2 = pd.DataFrame(columns=['a', 'b'], data=[[1, 2], [3, 4]]) - df3 = pd.DataFrame(columns=['c', 'd'], data=[[5, 6], [7, 8]]) + source_module='Preflight1', container=container_1) self._container_manager.StoreContainer( - source_module='Preflight1', - container=containers.DataFrame( - data_frame=df1, description='Description', name='name')) + source_module='Preflight1', container=container_2) self._container_manager.StoreContainer( - source_module='Preflight1', - container=containers.DataFrame( - data_frame=df2, description='Description', name='name')) + source_module='Preflight1', container=container_2) self._container_manager.StoreContainer( - source_module='Preflight1', - container=containers.DataFrame( - data_frame=df3, description='Description', name='name')) - - actual = self._container_manager.GetContainers( - requesting_module='ModuleA', container_class=_TestContainer1) - self.assertEqual(len(actual), 1) - self.assertIn(_TestContainer1('param1'), actual) - - actual = self._container_manager.GetContainers( - requesting_module='ModuleA', container_class=_TestContainer2) - self.assertEqual(len(actual), 2) - self.assertIn(_TestContainer2('param1'), actual) - self.assertIn(_TestContainer2('param2'), actual) - - actual = self._container_manager.GetContainers( - requesting_module='ModuleA', container_class=_TestContainer3) - self.assertEqual(len(actual), 1) - self.assertIn(_TestContainer3('param1'), actual) - - actual = self._container_manager.GetContainers( - requesting_module='ModuleA', container_class=containers.DataFrame) - self.assertEqual(len(actual), 2) - self.assertIn(containers.DataFrame( - data_frame=df1, description='Description', name='name'), actual) - self.assertIn(containers.DataFrame( - data_frame=df3, description='Description', name='name'), actual) + source_module='Preflight1', container=container_3) def test_ContainerStreaming(self): """Tests that container streaming operates as expected."""