Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Additional module testing refactoring. #951

Merged
merged 2 commits into from
Jan 15, 2025
Merged

Additional module testing refactoring. #951

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
290ba7b
Additional module testing refactoring.
ramo-j Jan 14, 2025
f1ceb1b
Linter appeasement
ramo-j Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion dftimewolf/lib/state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -564,7 +564,10 @@ def _RunModuleThread(self, module_definition: Dict[str, str]) -> None:
logger.info('Module {0:s} finished execution'.format(runtime_name))
self._threading_event_per_module[runtime_name].set()

self._container_manager.CompleteModule(runtime_name)
try:
self._container_manager.CompleteModule(runtime_name)
except Exception: # pylint: disable=broad-exception-caught
logger.warning('Unknown exception encountered', exc_info=True)

self.CleanUp()

Expand Down
2 changes: 1 addition & 1 deletion tests/lib/collectors/aws.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ def testProcess(self,
analysis_profile_name='test-analysis-profile-name',
analysis_zone='test-analysis-zone'
)
self._module.Process()
self._ProcessModule()

mock_CreateVolumeCopy.assert_called_with(
'test-remote-zone',
Expand Down
6 changes: 3 additions & 3 deletions tests/lib/collectors/aws_logging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def testProcess(self, mock_boto3):
query_filter='Username,fakename',
start_time=datetime.datetime(2021, 1, 1, 0, 0, 0),
end_time=datetime.datetime(2021, 1, 2, 0, 0, 0))
self._module.Process()
self._ProcessModule()

mock_session.client.assert_called_with(
'cloudtrail', region_name='fake-region')
Expand All @@ -80,13 +80,13 @@ def testProcess(self, mock_boto3):
mock_client.get_caller_identity.side_effect = (
boto_exceptions.NoCredentialsError)
with self.assertRaises(errors.DFTimewolfError):
self._module.Process()
self._ProcessModule()
mock_client.get_caller_identity.side_effect = None

mock_client.lookup_events.side_effect = (
boto_exceptions.ClientError({}, 'abc'))
with self.assertRaises(errors.DFTimewolfError):
self._module.Process()
self._ProcessModule()
mock_client.lookup_events.side_effect = None


Expand Down
10 changes: 2 additions & 8 deletions tests/lib/collectors/aws_snapshot_s3_copy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,10 +222,7 @@ def testProcessFromParams(self,
new=MockMakeAPICall):
self._module.SetUp(snaps_str, FAKE_BUCKET, FAKE_REGION)

self._module.PreProcess()
for c in self._module.GetContainers(containers.AWSSnapshot):
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

actual_output = [c.path for c in \
self._module.GetContainers(containers.AWSS3Object)]
Expand Down Expand Up @@ -268,10 +265,7 @@ def testProcessFromState(self,
new=MockMakeAPICall):
self._module.SetUp(None, FAKE_BUCKET, FAKE_REGION)

self._module.PreProcess()
for c in self._module.GetContainers(containers.AWSSnapshot):
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

actual_output = [c.path for c in \
self._module.GetContainers(containers.AWSS3Object)]
Expand Down
4 changes: 2 additions & 2 deletions tests/lib/collectors/aws_volume_snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def testProcessFromParams(self, mock_loader, mock_sleep):

with mock.patch('botocore.client.BaseClient._make_api_call',
new=MockMakeAPICall):
self._module.Process()
self._ProcessModule()

self.assertEqual(2, len(self._module.GetContainers(
containers.AWSSnapshot)))
Expand All @@ -123,7 +123,7 @@ def testProcessFromState(self, mock_loader, mock_sleep):

with mock.patch('botocore.client.BaseClient._make_api_call',
new=MockMakeAPICall):
self._module.Process()
self._ProcessModule()

self.assertEqual(2, len(self._module.GetContainers(
containers.AWSSnapshot)))
Expand Down
2 changes: 1 addition & 1 deletion tests/lib/collectors/azure.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ def testProcess(self,
analysis_region='test-analysis-region',
all_disks=True
)
self._module.Process()
self._ProcessModule()

mock_CreateDiskCopy.assert_called_with(
'test-analysis-resource-group-name',
Expand Down
6 changes: 3 additions & 3 deletions tests/lib/collectors/azure_logging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def testProcess(self, mock_monitor, mock_credentials):
self._module.SetUp(
subscription_id='55c5ff71-b3e2-450d-89da-cb12c1a38d87',
filter_expression='eventTimestamp ge \'2022-02-01\'')
self._module.Process()
self._ProcessModule()

mock_monitor.assert_called_with(
'Credentials', '55c5ff71-b3e2-450d-89da-cb12c1a38d87')
Expand All @@ -73,11 +73,11 @@ def testProcess(self, mock_monitor, mock_credentials):
# Ensure DFTimewolfError is raised when creds aren't found.
mock_credentials.side_effect = FileNotFoundError
with self.assertRaises(errors.DFTimewolfError):
self._module.Process()
self._ProcessModule()
mock_credentials.side_effect = None

# Ensure DFTimewolfError is raised when Azure libs raise an exception.
mock_activity_logs_client.list.side_effect = (
az_exceptions.HttpResponseError)
with self.assertRaises(errors.DFTimewolfError):
self._module.Process()
self._ProcessModule()
21 changes: 6 additions & 15 deletions tests/lib/collectors/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,8 @@ def testQuery(self, mock_bq):
"""Tests that the collector calls the BQ client."""
mock_bq().query().to_dataframe().to_json.return_value = "{'foo':1}"
self._module.SetUp('test_project', 'test_query', 'test_description', False)
self._module.PreProcess()
for c in self._module.GetContainers(
self._module.GetThreadOnContainerType()):
self._module.Process(c) # pytype: disable=wrong-arg-types
self._module.PostProcess()
self._ProcessModule()

mock_bq().query.assert_called_with('test_query')
mock_bq().query().to_dataframe().to_json.assert_called_once()

Expand All @@ -45,11 +42,8 @@ def testQueryFromState(self, mock_bq):
cont_in.SetMetadata('input_metadata_key', 'input_metadata_value')
self._module.StoreContainer(cont_in)
self._module.SetUp('test_project', '', '', False)
self._module.PreProcess()
for c in self._module.GetContainers(
self._module.GetThreadOnContainerType()):
self._module.Process(c) # pytype: disable=wrong-arg-types
self._module.PostProcess()
self._ProcessModule()

mock_bq().query.assert_called_with('test_query')

conts = self._module.GetContainers(containers.DataFrame)
Expand All @@ -62,11 +56,8 @@ def testQueryPandaOutput(self, mock_bq):
"""Tests placing query results in a dataframe."""
mock_bq().query().to_dataframe.return_value = pd.DataFrame([1], ['foo'])
self._module.SetUp('test_project', 'test_query', 'test_description', True)
self._module.PreProcess()
for c in self._module.GetContainers(
self._module.GetThreadOnContainerType()):
self._module.Process(c) # pytype: disable=wrong-arg-types
self._module.PostProcess()
self._ProcessModule()

mock_bq().query.assert_called_with('test_query')

conts = self._module.GetContainers(containers.DataFrame)
Expand Down
2 changes: 1 addition & 1 deletion tests/lib/collectors/filesystem.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def testOutput(self, mock_exists):
fake_paths = '/fake/path/1,/fake/path/2'
self._module.SetUp(paths=fake_paths)
mock_exists.return_value = True
self._module.Process()
self._ProcessModule()
files = self._module.GetContainers(containers.File)
self.assertEqual(files[0].path, '/fake/path/1')
self.assertEqual(files[0].name, '1')
Expand Down
92 changes: 13 additions & 79 deletions tests/lib/collectors/grr_hosts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,12 +245,7 @@ def testProcessSpecificArtifacts(self,
skip_offline_clients=False
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

kwargs = mock_ArtifactCollectorFlowArgs.call_args[1]
self.assertTrue(kwargs['ignore_interpolation_errors']) # default argument
Expand Down Expand Up @@ -291,12 +286,7 @@ def testProcessWindowsArtifacts(self,
skip_offline_clients=False
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

kwargs = mock_ArtifactCollectorFlowArgs.call_args[1]
self.assertFalse(kwargs["use_raw_filesystem_access"])
Expand Down Expand Up @@ -333,12 +323,7 @@ def testProcessLinuxArtifacts(self,
skip_offline_clients=False
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

kwargs = mock_ArtifactCollectorFlowArgs.call_args[1]
# No raw access for Linux
Expand All @@ -355,12 +340,7 @@ def testProcess(self, mock_CreateFlow, mock_Get, mock_DownloadFiles):
mock_DownloadFiles.return_value = '/tmp/tmpRandom/tomchop'
mock_Get.return_value = mock_grr_hosts.MOCK_FLOW

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

# Flow ID is F:12345, Client ID is C.0000000000000001
self.mock_grr_api.SearchClients.assert_any_call('C.0000000000000001')
Expand Down Expand Up @@ -407,12 +387,7 @@ def testProcessFromHostContainers(
self._module.StoreContainer(
containers.Host(hostname='container.host'))

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_FindClients.assert_called_with(['container.host'])

Expand Down Expand Up @@ -455,13 +430,7 @@ def testProcessFromArtifactContainers(
containers.GRRArtifact(name="ForensicArtifact2")
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType()
)
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

actual_flow_proto = mock_LaunchFlow.call_args_list[0][0][2]
self.assertListEqual(
Expand Down Expand Up @@ -560,12 +529,7 @@ def testProcess(self, mock_LaunchFlow, mock_DownloadFiles, _, mock_InitHttp):
skip_offline_clients=False,
action='stat',
)
self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_LaunchFlow.assert_called_with(
mock_grr_hosts.MOCK_CLIENT_RECENT,
Expand Down Expand Up @@ -615,12 +579,7 @@ def testWindowsProcess(
action='stat',
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_LaunchFlow.assert_called_with(
mock_grr_hosts.MOCK_WINDOWS_CLIENT,
Expand Down Expand Up @@ -670,12 +629,7 @@ def testProcessFromContainers(self,
self._module.StoreContainer(
containers.Host(hostname='container.host'))

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_FindClients.assert_called_with(['container.host'])

Expand Down Expand Up @@ -723,12 +677,7 @@ def testProcess(self, _, mock_DownloadFiles, unused_mock_verify_access):
mock_grr_hosts.MOCK_CLIENT_LIST
mock_DownloadFiles.return_value = '/tmp/something'

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_DownloadFiles.assert_called_once_with(
mock_grr_hosts.MOCK_CLIENT_RECENT, 'F:12345')
Expand Down Expand Up @@ -808,12 +757,7 @@ def testProcessNoFlowData(
approvers='approver1,approver2',
skip_offline_clients=False,
)
self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

log_messages = [record.getMessage() for record in lc.records]
# pylint: disable=line-too-long
Expand Down Expand Up @@ -874,12 +818,7 @@ def testProcess(self):
mock_createflow.return_value.flow_id = "F:12345"
mock_getcollectedtimeline.return_value.WriteToFile = _MOCK_WRITE_TO_FILE

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c) # pytype: disable=wrong-arg-count
self._module.PostProcess()
self._ProcessModule()

mock_createflow.assert_called_once_with(
name="TimelineFlow", args=timeline_pb2.TimelineArgs(root=b"/")
Expand Down Expand Up @@ -920,12 +859,7 @@ def testProcessFromContainers(
skip_offline_clients=False,
)

self._module.PreProcess()
in_containers = self._module.GetContainers(
self._module.GetThreadOnContainerType())
for c in in_containers:
self._module.Process(c)
self._module.PostProcess()
self._ProcessModule()

mock_FindClients.assert_called_with(['container.host'])

Expand Down
Loading
Loading