diff --git a/src/brave/types/web/web_search_response.py b/src/brave/types/web/web_search_response.py index 96a4e9c..0c7bf74 100644 --- a/src/brave/types/web/web_search_response.py +++ b/src/brave/types/web/web_search_response.py @@ -89,9 +89,13 @@ def video_results(self) -> List[str]: return _results.get("videos").get("results") if self.videos else [] @property - def product_cluster(self) -> List[str]: + def product_cluster(self) -> List: """Return a list of product clusters.""" - return [result.product_cluster for result in self._web_results if result.subtype == "product_cluster"][0] + clusters = [] + for result in self._web_results: + if result.subtype == "product_cluster" and result.product_cluster: + clusters.extend(result.product_cluster) + return clusters def download_all_pdfs(self, path: str = "downloads") -> None: """Download PDFs for all search results.""" @@ -101,7 +105,7 @@ def download_all_pdfs(self, path: str = "downloads") -> None: def product_prices(self) -> List[int]: """Return a list of product prices.""" - if len(self.product_cluster) == 0: + if not self.product_cluster: return [ float(result.product.price) for result in self._web_results if result.product and result.product.price ] @@ -111,11 +115,13 @@ def product_prices(self) -> List[int]: def product_price_ranges(self) -> Tuple[int, int]: """Return a list of product price ranges.""" prices = self.product_prices() + if not prices: + return (0, 0) return (min(prices), max(prices)) def average_product_review_score(self) -> Optional[float]: """Return the average product review score out of 100.""" - if len(self.product_cluster) == 0: + if not self.product_cluster: ratings = [ float(result.product.rating.ratingValue) / float(result.product.rating.bestRating) for result in self._web_results diff --git a/tests/unit/test_web_search_response.py b/tests/unit/test_web_search_response.py new file mode 100644 index 0000000..1f81c1b --- /dev/null +++ b/tests/unit/test_web_search_response.py @@ -0,0 +1,29 @@ +import json + +from brave.types import WebSearchApiResponse + + +def _load_response(): + with open("tests/test_responses/blue_tack_minimal.json", "r") as f: + data = json.load(f) + return WebSearchApiResponse.model_validate(data) + + +def test_product_cluster_empty(): + response = _load_response() + assert response.product_cluster == [] + + +def test_product_prices_empty(): + response = _load_response() + assert response.product_prices() == [] + + +def test_product_price_ranges_empty(): + response = _load_response() + assert response.product_price_ranges() == (0, 0) + + +def test_average_product_review_score_none(): + response = _load_response() + assert response.average_product_review_score() is None