def ListGrrBinaries(self): self.request_count += 1 if self._raise_conn_error and self.request_count == 1: raise requests.ConnectionError("Fake Connection Error.") linux_binary = mock.Mock(spec=api_config.GrrBinary) linux_binary.path = runner.E2ETestRunner.LINUX_TEST_BINARY_PATH windows_binary = mock.Mock(spec=api_config.GrrBinary) windows_binary.path = runner.E2ETestRunner.WINDOWS_TEST_BINARY_PATH return api_utils.ItemsIterator(items=[linux_binary, windows_binary])
def SendIteratorRequest(self, handler_name, args): if not args or not hasattr(args, "count"): result = self.connector.SendRequest(handler_name, args) total_count = getattr(result, "total_count", None) return utils.ItemsIterator(items=result.items, total_count=total_count) else: pages = self._GeneratePages(handler_name, args) first_page = pages.next() total_count = getattr(first_page, "total_count", None) page_items = lambda page: page.items next_pages_items = itertools.chain.from_iterable(map(page_items, pages)) all_items = itertools.chain(first_page.items, next_pages_items) if args.count: all_items = itertools.islice(all_items, args.count) return utils.ItemsIterator(items=all_items, total_count=total_count)
def SendIteratorRequest( self, handler_name: str, args: Any, ) -> utils.ItemsIterator: """Sends an iterator request.""" if not args or not hasattr(args, "count"): result = self.connector.SendRequest(handler_name, args) if not hasattr(result, "items"): detail = f"Incorrect result type for '{handler_name}': {type(result)}" raise TypeError(detail) total_count = getattr(result, "total_count", None) return utils.ItemsIterator(items=result.items, total_count=total_count) else: pages = self._GeneratePages(handler_name, args) first_page = next(pages) total_count = getattr(first_page, "total_count", None) def PageItems(page: message.Message) -> Iterator[message.Message]: if not hasattr(page, "items"): detail = f"Incorrect page type for '{handler_name}': {type(page)}" raise TypeError(detail) return page.items next_pages_items = itertools.chain.from_iterable( map(PageItems, pages)) all_items = itertools.chain(PageItems(first_page), next_pages_items) if args.count: all_items = itertools.islice(all_items, args.count) return utils.ItemsIterator(items=all_items, total_count=total_count)