Async Query Client

You are given the class QueryClient below. Its API methods are already implemented and scripted by the test; you only implement run.

The tests build the client with the status responses to return in order — the last one repeats once the script runs out — and a dict of pages keyed by token, call run, then call stats.

class QueryClient:
    def __init__(self, polls: list[dict], pages: dict[str, dict], max_polls: int) -> None:
        self.max_polls = max_polls
        self._polls = polls
        self._pages = pages
        self._submits = 0
        self._status_calls = 0
        self._page_calls = 0

    def submit_query(self, sql: str) -> dict:
        """Starts a query and returns {"query_id": ...}."""
        self._submits += 1
        return {"query_id": f"q{self._submits}"}

    def get_query_status(self, query_id: str) -> dict:
        """The next scripted response; the last one repeats once the script runs out."""
        index = min(self._status_calls, len(self._polls) - 1)
        self._status_calls += 1
        return self._polls[index]

    def fetch_result_page(self, query_id: str, token: str) -> dict:
        """The page stored under `token`: {"rows": [...], "next_token": ... or None}."""
        self._page_calls += 1
        return self._pages[token]

    def stats(self) -> list[int]:
        """[queries submitted, status calls made, page fetches made]."""
        return [self._submits, self._status_calls, self._page_calls]

    def run(self, sql: str) -> list:
        raise NotImplementedError