def test_reconstruct_all_subresponses(self):
# Construct a fake stoq_response as if it were generated from a file
# A.zip that contains two files, B.txt and C.zip, where C.zip contains D.txt
initial_response = StoqResponse(
results=[
PayloadResults(
payload_id="A.zip",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result1"}],
plugins_run={"workers": [["fake"]]},
),
PayloadResults(
payload_id="B.txt",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result2"}],
plugins_run={"workers": [["fake"]]},
extracted_from="A.zip",
extracted_by="fake",
),
PayloadResults(
payload_id="C.zip",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result3"}],
plugins_run={"workers": [["fake"]]},
extracted_from="A.zip",
extracted_by="fake",
),
PayloadResults(
payload_id="D.txt",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result4"}],
plugins_run={"workers": [["fake"]]},
extracted_from="C.zip",
extracted_by="fake",
),
],
request_meta=RequestMeta(extra_data={"check": "me"}),
errors={},
)
s = Stoq(base_dir=utils.get_data_dir(), decorators=["simple_decorator"])
all_subresponses = list(s.reconstruct_all_subresponses(initial_response))
# We expect there to be four "artificial" responses generated, one for
# each payload as the root.
self.assertEqual(len(all_subresponses), 4)
# We expect the first response to have all 4 payloads, the second response
# to have just the second payload, the third response to have the third
# and fourth payload, and the fourth response to have just the fourth payload
self.assertEqual(
[len(stoq_response.results) for stoq_response in all_subresponses], [4, 1, 2, 1]
)
self.assertEqual(
[
stoq_response.results[0].workers[0]["fake"]
for stoq_response in all_subresponses
],
["result1", "result2", "result3", "result4"],
)
self.assertTrue(
all(
"simple_decorator" in stoq_response.decorators
for stoq_response in all_subresponses
)
)
# Assert that they all have the same scan ID
self.assertEqual(
len({stoq_response.scan_id for stoq_response in all_subresponses}), 1
)
async def test_reconstruct_all_subresponses(self):
# Construct a fake stoq_response as if it were generated from a file
# A.zip that contains two files, B.txt and C.zip, where C.zip contains D.txt
results = [
Payload(content=b'', payload_id='A.zip', payload_meta=PayloadMeta()),
Payload(
content=b'',
payload_id='B.txt',
payload_meta=PayloadMeta(),
extracted_from='A.zip',
extracted_by='fake',
),
Payload(
content=b'',
payload_id='C.zip',
payload_meta=PayloadMeta(),
extracted_from='A.zip',
extracted_by='fake',
),
Payload(
content=b'',
payload_id='D.txt',
payload_meta=PayloadMeta(),
extracted_from='C.zip',
extracted_by='fake',
),
]
request = Request(request_meta=RequestMeta(extra_data={'check': 'me'}))
payload_count = 1
for result in results:
result.results.workers['fake'] = f'result-{payload_count}'
result.results.plugins_run['workers'].append('fake')
request.payloads.append(result)
payload_count += 1
initial_response = StoqResponse(request)
s = Stoq(base_dir=utils.get_data_dir(), decorators=['simple_decorator'])
all_subresponses = [
r async for r in s.reconstruct_all_subresponses(initial_response)
]
# We expect there to be four "artificial" responses generated, one for
# each payload as the root.
self.assertEqual(len(all_subresponses), 4)
# We expect the first response to have all 4 payloads, the second response
# to have just the second payload, the third response to have the third
# and fourth payload, and the fourth response to have just the fourth payload
self.assertEqual(
[len(stoq_response.results) for stoq_response in all_subresponses],
[4, 1, 2, 1],
)
self.assertEqual(
[
stoq_response.results[0].workers['fake']
for stoq_response in all_subresponses
],
['result-1', 'result-2', 'result-3', 'result-4'],
)
self.assertTrue(
all(
'simple_decorator' in stoq_response.decorators
for stoq_response in all_subresponses
)
)
# Assert that they all have the same scan ID
self.assertEqual(
len({stoq_response.scan_id for stoq_response in all_subresponses}), 1
)
