From 5e0fa1c54bd37155a5301d608bd50c8e4dd5ddc3 Mon Sep 17 00:00:00 2001
From: "tiffany.cheng" <tiffany.cheng@jisc.ac.uk>
Date: Thu, 18 Jun 2026 14:30:55 +0100
Subject: [PATCH] Fix: use field mapping in create_total_query

---
 mario/query_builder.py        |  1 +
 test/test_athena_extractor.py | 48 +++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/mario/query_builder.py b/mario/query_builder.py
index c5110c9..780e45c 100644
--- a/mario/query_builder.py
+++ b/mario/query_builder.py
@@ -108,6 +108,7 @@ def create_totals_query(self, measure=None) -> [str, List[any]]:
         TODO make this more generic
         """
         measures = []
+        measure = self.mapping.as_physical[measure]
         if measure is None:
             for measure in self.dataset_specification.measures:
                 # Decide column name based on whether 'subject' fields are present
diff --git a/test/test_athena_extractor.py b/test/test_athena_extractor.py
index 248903b..a54db1b 100644
--- a/test/test_athena_extractor.py
+++ b/test/test_athena_extractor.py
@@ -56,6 +56,37 @@ def get_test_conf():
 
     return dataset, metadata, config
 
+def get_test_conf_tdsa():
+    dataset = DatasetSpecification()
+    dataset.dimensions = [
+        "Academic Year",
+        "Mode of study",
+        "Country of HE provider"
+    ]
+    dataset.measures = ['Unrounded FPE']
+    dataset.name = 'v_sp_student_fpe'
+    metadata = Metadata()
+    academic_year = Item()
+    academic_year.name = 'Academic Year'
+    mode_of_study = Item()
+    mode_of_study.name = 'Mode of study'
+    mode_of_study.set_property('domain', ['Full-time', 'Part-time'])
+    country_of_he_provider = Item()
+    country_of_he_provider.name = 'Country of HE provider'
+    country_of_he_provider.set_property('domain', ['England', 'Wales', 'Scotland', 'Northern Ireland'])
+    fpe_field = Item()
+    fpe_field.name = 'Unrounded FPE'
+    metadata.add_item(academic_year)
+    metadata.add_item(mode_of_study)
+    metadata.add_item(fpe_field)
+    metadata.add_item(country_of_he_provider)
+
+    config = AthenaConfiguration()
+    config.aws_s3_staging_dir = AWS_ATHENA_RESULTS_DIR
+    config.aws_region_name = AWS_REGION
+
+    return dataset, metadata, config
+
 
 def test_athena_count():
     if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION:
@@ -75,6 +106,23 @@ def test_athena_count():
     total = extractor.get_total(measure=dataset.measures[0])
     assert total == 28_733_910
 
+def test_athena_count_with_space():
+    if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION:
+        pytest.skip("Skipping Athena test as AWS not configured")
+
+    dataset, metadata, cfg = get_test_conf_tdsa()
+    cfg.query_builder = SubsetQueryBuilder
+    cfg.schema = 'dev'
+    cfg.view = 'v_sp_student_fpe'
+
+    extractor = AthenaDataExtractor(
+        configuration=cfg,
+        metadata=metadata,
+        dataset_specification=dataset
+    )
+
+    total = extractor.get_total(measure=dataset.measures[0])
+    assert total == 3_332_785
 
 def test_athena_save_data_as_csv():
     if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: