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 )