diff --git a/jobs/quicksuggest2bq/quicksuggest2bq/main.py b/jobs/quicksuggest2bq/quicksuggest2bq/main.py index 8f63c89e..f5b84960 100644 --- a/jobs/quicksuggest2bq/quicksuggest2bq/main.py +++ b/jobs/quicksuggest2bq/quicksuggest2bq/main.py @@ -101,6 +101,10 @@ def download_suggestions(client: kinto_http.Client) -> Iterator[KintoSuggestion] {"keyword": kw, "count": count} for kw, count in suggestion_data.get("full_keywords", []) ], + "serp_categories": [ + {"category": category_id} + for category_id in suggestion_data.get("serp_categories", []) + ], } yield KintoSuggestion(**suggestion) diff --git a/jobs/quicksuggest2bq/tests/test_main.py b/jobs/quicksuggest2bq/tests/test_main.py index 8ae29bfb..185a9535 100644 --- a/jobs/quicksuggest2bq/tests/test_main.py +++ b/jobs/quicksuggest2bq/tests/test_main.py @@ -135,7 +135,8 @@ def test_suggestion_full_keyword(self, mocked_kinto_client): def test_suggestion_serp_categories(self, mocked_kinto_client): suggestions = list(download_suggestions(mocked_kinto_client)) - assert suggestions[0].serp_categories == SERP_CATEGORIES + assert suggestions[0].serp_categories[0]["category"] == 1 + assert suggestions[0].serp_categories[1]["category"] == 2 def test_suggestion_score(self, mocked_kinto_client): suggestions = list(download_suggestions(mocked_kinto_client))