From bf8c7a2f224bcd6f9e679f5670187b3f9a3a003c Mon Sep 17 00:00:00 2001 From: munnsmunns Date: Wed, 3 Jan 2024 15:18:30 -0500 Subject: [PATCH] raise HawcClientException instead of timing out --- client/hawc_client/client.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/client/hawc_client/client.py b/client/hawc_client/client.py index aa464875d6..4e6f470f9c 100644 --- a/client/hawc_client/client.py +++ b/client/hawc_client/client.py @@ -7,6 +7,7 @@ from playwright.sync_api import Page from playwright.sync_api._context_manager import PlaywrightContextManager as pcm +from .exceptions import HawcClientException from .session import HawcSession @@ -128,7 +129,10 @@ def download_visual(self, id: int, is_tableau: bool = False, fn: PathLike = None BytesIO: the PNG representation of the visual, in bytes. """ url = f"{self.client.session.root_url}/summary/visual/{id}/" - self.page.goto(url) + # ensure response is OK before waiting + response = self.page.goto(url) + if response and not response.ok: + raise HawcClientException(response.status, response.status_text) data = fetch_png(self.page, is_tableau) write_to_file(data, fn) @@ -146,7 +150,10 @@ def download_data_pivot(self, id: int, fn: PathLike = None) -> BytesIO: BytesIO: the PNG representation of the data pivot, in bytes. """ url = f"{self.client.session.root_url}/summary/data-pivot/{id}/" - self.page.goto(url) + # ensure response is OK before waiting + response = self.page.goto(url) + if response and not response.ok: + raise HawcClientException(response.status, response.status_text) data = fetch_png(self.page) write_to_file(data, fn) return data