From c1719df4d996e5d8e21178675fa0908286bbed8f Mon Sep 17 00:00:00 2001
From: Casper Welzel Andersen <casper.andersen@epfl.ch>
Date: Tue, 8 Jun 2021 19:14:41 +0200
Subject: [PATCH] Fix some tests

Make all parts passed as mongo_kwargs to `update_query` clean Python
types.
---
 optimade_gateway/queries/utils.py   | 2 +-
 optimade_gateway/routers/queries.py | 5 +++--
 tests/routers/test_info.py          | 3 ++-
 tests/routers/test_queries.py       | 9 ++++++---
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/optimade_gateway/queries/utils.py b/optimade_gateway/queries/utils.py
index 5e3903bb..21656b31 100644
--- a/optimade_gateway/queries/utils.py
+++ b/optimade_gateway/queries/utils.py
@@ -47,7 +47,7 @@ async def update_query(
     update_kwargs = {"$set": {"last_modified": update_time}}
 
     if mongo_kwargs:
-        update_kwargs.update(mongo_kwargs)
+        update_kwargs.update(await clean_python_types(mongo_kwargs))
 
     if operator and operator == "$set":
         update_kwargs["$set"].update({field: await clean_python_types(value)})
diff --git a/optimade_gateway/routers/queries.py b/optimade_gateway/routers/queries.py
index ba43bc96..0ce8867c 100644
--- a/optimade_gateway/routers/queries.py
+++ b/optimade_gateway/routers/queries.py
@@ -25,6 +25,7 @@
 from optimade_gateway.common.config import CONFIG
 from optimade_gateway.mappers import QueryMapper
 from optimade_gateway.models import (
+    GatewayQueryResponse,
     QueryCreate,
     QueryResource,
     QueryState,
@@ -114,7 +115,7 @@ async def post_queries(
 
 @ROUTER.get(
     "/queries/{query_id:path}",
-    response_model=Union[EntryResponseMany, ErrorResponse],
+    response_model=Union[EntryResponseMany, ErrorResponse, GatewayQueryResponse],
     response_model_exclude_defaults=False,
     response_model_exclude_none=False,
     response_model_exclude_unset=True,
@@ -124,7 +125,7 @@ async def get_query(
     request: Request,
     query_id: str,
     response: Response,
-) -> Union[EntryResponseMany, ErrorResponse]:
+) -> Union[EntryResponseMany, ErrorResponse, GatewayQueryResponse]:
     """`GET /queries/{query_id}`
 
     Return the response from a query
diff --git a/tests/routers/test_info.py b/tests/routers/test_info.py
index 597687bd..2f7512e4 100644
--- a/tests/routers/test_info.py
+++ b/tests/routers/test_info.py
@@ -33,9 +33,10 @@ async def test_get_info(
             "redoc",
             "search",
             "queries",
+            "databases",
         ]
     )
-    entry_types_by_format = {"json": ["gateways", "queries"]}
+    entry_types_by_format = {"json": ["databases", "gateways", "queries"]}
 
     response = await client("/info")
 
diff --git a/tests/routers/test_queries.py b/tests/routers/test_queries.py
index 57d55b5b..d3b15366 100644
--- a/tests/routers/test_queries.py
+++ b/tests/routers/test_queries.py
@@ -156,10 +156,13 @@ async def test_query_results(
     """Test POST /queries and GET /queries/{id}"""
     import asyncio
     from optimade.models import EntryResponseMany
-    from optimade.models import StructureResponseMany
 
     from optimade_gateway.common.config import CONFIG
-    from optimade_gateway.models.queries import QueryState, QueryResource
+    from optimade_gateway.models.queries import (
+        GatewayQueryResponse,
+        QueryState,
+        QueryResource,
+    )
 
     data = {
         "id": "test",
@@ -193,7 +196,7 @@ async def test_query_results(
     response = await client(f"/queries/{data['id']}")
     assert response.status_code == 200, f"Request failed: {response.json()}"
 
-    response = StructureResponseMany(**response.json())
+    response = GatewayQueryResponse(**response.json())
     assert response.data
     assert (
         getattr(response.meta, f"_{CONFIG.provider.prefix}_query", "NOT FOUND")