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)
