From f76263fe5b9e424337ec9a9579a2a28dc3c4a425 Mon Sep 17 00:00:00 2001
From: olaughter <fred@climatepolicyradar.org>
Date: Wed, 18 Sep 2024 11:43:16 +0100
Subject: [PATCH] Ensure geography fields are returned

---
 .../test_documents/family_document.json       |  2 +-
 tests/test_search_adaptors.py                 | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/tests/local_vespa/test_documents/family_document.json b/tests/local_vespa/test_documents/family_document.json
index 10e0df12..644c605c 100644
--- a/tests/local_vespa/test_documents/family_document.json
+++ b/tests/local_vespa/test_documents/family_document.json
@@ -309,7 +309,7 @@
       "family_slug": "environmental-strategy-for-2014-2023_9f8e",
       "document_source_url": "https://wedocs.unep.org/bitstream/handle/20.500.11822/9507/-Environmental_Strategy_for_the_years_2014-2023-2014Moldova_EnvironmentalStrategy_2014-202.pdf?sequence=3&isAllowed=y",
       "family_geography": "MDA",
-      "family_geographies": null,
+      "family_geographies": ["MDA"],
       "family_category": "Executive",
       "document_md5_sum": "bea7a05dae73fbbd629e687a71a15b95",
       "family_name_index": "Environmental Strategy for 2014-2023",
diff --git a/tests/test_search_adaptors.py b/tests/test_search_adaptors.py
index ffc85111..4505914d 100644
--- a/tests/test_search_adaptors.py
+++ b/tests/test_search_adaptors.py
@@ -7,6 +7,7 @@
 from cpr_sdk.models.search import (
     Document,
     Filters,
+    Hit,
     MetadataFilter,
     Passage,
     SearchParameters,
@@ -14,6 +15,8 @@
     sort_fields,
 )
 from cpr_sdk.search_adaptors import VespaSearchAdapter
+from cpr_sdk.utils import dig
+from cpr_sdk.vespa import build_vespa_request_body
 
 
 def vespa_search(
@@ -554,3 +557,39 @@ def test_vespa_search_adaptor__filters(test_vespa, query_string, filters):
                 attribute_value_from_hit = getattr(hit, filter_name)
                 assert attribute_value_from_hit not in [None, []]
                 assert all([val in attribute_value_from_hit for val in filter_values])
+
+
+@pytest.mark.vespa
+@pytest.mark.parametrize("query_string", ["e"])
+@pytest.mark.parametrize("exact_match", [True, False])
+@pytest.mark.parametrize(
+    "metadata_filters", [None, [{"name": "family.sector", "value": "Price"}]]
+)
+@pytest.mark.parametrize("geographies", [None, {"family_geographies": ["BIH"]}])
+def test_vespa_search_response__geographies(
+    test_vespa, query_string, exact_match, metadata_filters, geographies
+) -> None:
+    """Test that the search response includes geographies"""
+    parameters = SearchParameters(
+        query_string=query_string,
+        exact_match=exact_match,
+        filters=Filters.model_validate(geographies) if geographies else None,
+        metadata=(
+            [
+                MetadataFilter.model_validate(metadata_filter)
+                for metadata_filter in metadata_filters
+            ]
+            if metadata_filters
+            else None
+        ),
+    )
+
+    vespa_response = test_vespa.client.query(body=build_vespa_request_body(parameters))
+
+    root = vespa_response.json["root"]
+    response_families = dig(root, "children", 0, "children", 0, "children", default=[])
+    for family in response_families:
+        for hit in dig(family, "children", 0, "children", default=[]):
+            hit = Hit.from_vespa_response(response_hit=hit)
+            assert hit.family_geography not in [None, []]
+            assert hit.family_geographies not in [None, []]