From 44b707a73a415901b30fa8205f780d8c40b06577 Mon Sep 17 00:00:00 2001
From: Matthew Richards <matthew.richards@stfc.ac.uk>
Date: Tue, 13 Apr 2021 13:04:37 +0000
Subject: [PATCH]  #141: Fix bug when a count endpoint with distinct filter
 should get 0 results but caused an IndexError

- Test case to cover this edge case also added in this commit
---
 datagateway_api/common/icat/query.py | 11 ++++++-----
 test/icat/test_query.py              | 17 ++++++++++++++---
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/datagateway_api/common/icat/query.py b/datagateway_api/common/icat/query.py
index 284305a8..5c20649d 100644
--- a/datagateway_api/common/icat/query.py
+++ b/datagateway_api/common/icat/query.py
@@ -116,12 +116,13 @@ def execute_query(self, client, return_json_formattable=False):
             log.info("Query results will be returned in a JSON format")
             data = []
 
+            if self.query.manual_count:
+                # Manually count the number of results
+                data.append(len(query_result))
+                return data
+
             for result in query_result:
-                if self.query.manual_count:
-                    # Manually count the number of results
-                    data.append(len(query_result))
-                    break
-                elif distinct_query:
+                if distinct_query:
                     # When multiple attributes are given in a distinct filter, Python
                     # ICAT returns the results in a nested list. This doesn't happen
                     # when a single attribute is given, so the result is encased in a
diff --git a/test/icat/test_query.py b/test/icat/test_query.py
index 7965d07b..b1a38d0d 100644
--- a/test/icat/test_query.py
+++ b/test/icat/test_query.py
@@ -4,7 +4,7 @@
 import pytest
 
 from datagateway_api.common.date_handler import DateHandler
-from datagateway_api.common.exceptions import FilterError, PythonICATError
+from datagateway_api.common.exceptions import PythonICATError
 from datagateway_api.common.icat.filters import (
     PythonICATSkipFilter,
     PythonICATWhereFilter,
@@ -12,7 +12,7 @@
 from datagateway_api.common.icat.query import ICATQuery
 
 
-def prepare_icat_data_for_assertion(data, remove_id=False):
+def prepare_icat_data_for_assertion(data, remove_id=False, remove_visit_id=False):
     """
     Remove meta attributes from ICAT data. Meta attributes contain data about data
     creation/modification, and should be removed to ensure correct assertion values
@@ -38,6 +38,8 @@ def prepare_icat_data_for_assertion(data, remove_id=False):
         # meta_attributes is immutable
         if remove_id:
             entity.pop("id")
+        if remove_visit_id:
+            entity.pop("visitId")
 
         assertable_data.append(entity)
 
@@ -172,7 +174,6 @@ def test_invalid_query_creation(self, icat_client):
                 ],
                 id="Query with included entity",
             ),
-            # pytest.param(id="Query with included entity"),  # facility?
             pytest.param(
                 {
                     "title": "like '%Test data for the Python ICAT Backend on"
@@ -260,6 +261,16 @@ def test_invalid_query_creation(self, icat_client):
                 [1],
                 id="Multiple distinct fields on count query",
             ),
+            pytest.param(
+                {"title": "like '%Unknown testing data for DG API%'",},
+                "DISTINCT",
+                None,
+                ["title", "name"],
+                True,
+                True,
+                [0],
+                id="Multiple distinct fields on count query to return 0 matches",
+            ),
         ],
     )
     @pytest.mark.usefixtures("single_investigation_test_data")