Exemple #1
0
    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
        )