Mock the AWS CostExplorer API + split some methods

- The methods on the AWSCostExplorer class should be doing one
  thing at a time, and respond with data specific to it, rather than
  anything to do with the infinity data source. This allows us to
  test them better, as well as re-use this code in other ways
  such as 2i2c-org/initiatives#55
- Mock the HTTP server that's responding to CE requests, and check-in
  a mock response that we can use to validate our code - a true
  integration test.
- Allow configuring the AWS CE client we are using via traitlets. We use
  this here for tests, but we can also use this in other places where
  we want to configure other params.
- Split the AWS Integration tests into a file by themselves

Author: yuvipanda

src/jupyterhub_cost_monitoring/app.py

@@ -80,7 +80,15 @@ def total_costs(
     date_range = parse_from_to_in_query_params(from_date, to_date)
 
     try:
-        return aws_ce.query_total_costs(date_range)
+        account_costs = aws_ce.query_account_costs(date_range)
+        attributable_costs = aws_ce.query_attributable_costs(date_range)
+
+        # the infinity plugin appears needs us to sort by date, otherwise it fails
+        # to distinguish time series by the name field for some reason
+        sorted_response = sorted(
+            account_costs + attributable_costs, key=lambda x: x["date"]
+        )
+        return sorted_response
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"{e}")
 

src/jupyterhub_cost_monitoring/aws.py

@@ -8,7 +8,7 @@
 
 import boto3
 import requests
-from traitlets import Instance
+from traitlets import Dict, Instance, Unicode
 from traitlets.config import LoggingConfigurable
 
 from .cache import ttl_lru_cache
@@ -43,15 +43,25 @@ def _get_component_name(self, service_name):
         klass=Prometheus,
     )
 
+    aws_client_extra_kwargs = Dict(
+        Unicode(),
+        help="""
+        Extra arguments to be passed to the AWS Client that talks to the Cost Explorer
+        """,
+        config=True,
+    )
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.aws_ce_client = boto3.client("ce")
+        self.aws_ce_client = boto3.client("ce", **self.aws_client_extra_kwargs)
 
-    def query(self, metrics, granularity, from_date, to_date, filter, group_by):
+    def query(self, metrics, granularity, date_range: DateRange, filter, group_by):
         """
         Function meant to be responsible for making the API call and handling
         pagination etc. Currently pagination isn't handled.
         """
+        from_date, to_date = date_range.aws_range
+
         # ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ce/client/get_cost_and_usage.html#get-cost-and-usage
         response = self.aws_ce_client.get_cost_and_usage(
             Metrics=metrics,
@@ -93,84 +103,51 @@ def query_hub_names(self, date_range: DateRange):
         hub_names = [t or "support" for t in response["Tags"]]
         return hub_names
 
-    @ttl_lru_cache(seconds_to_live=3600)
-    def query_total_costs(self, date_range: DateRange):
-        """
-        Query total costs from AWS Cost Explorer for the given date range.
-
-        Reports both the total AWS account cost and the total attributable cost.
-        Not all costs will be successfully attributed, such as the cost of accessing
-        the AWS Cost Explorer API - its not something that can be attributed based
-        on a tag.
-
-        Args:
-            date_range: DateRange object containing the time period for the query
+    def query_account_costs(self, date_range: DateRange):
+        from_date, to_date = date_range.aws_range
 
-        Returns:
-            List of cost entries with 'date', 'cost', and 'name' fields, sorted by date
-        """
-        total_account_costs = self._query_total_costs(
-            date_range, add_attributable_costs_filter=False
-        )
-        total_attributable_costs = self._query_total_costs(
-            date_range, add_attributable_costs_filter=True
+        response = self.query(
+            metrics=[METRICS_UNBLENDED_COST],
+            granularity=GRANULARITY_DAILY,
+            date_range=date_range,
+            filter=FILTER_USAGE_COSTS,
+            group_by=[],
         )
 
-        processed_response = total_account_costs + total_attributable_costs
-
-        # the infinity plugin appears needs us to sort by date, otherwise it fails
-        # to distinguish time series by the name field for some reason
-        processed_response = sorted(processed_response, key=lambda x: x["date"])
+        processed_response = [
+            {
+                "date": e["TimePeriod"]["Start"],
+                "cost": f"{float(e['Total']['UnblendedCost']['Amount']):.2f}",
+                "name": "account",
+            }
+            for e in response["ResultsByTime"]
+        ]
 
         return processed_response
 
-    @ttl_lru_cache(seconds_to_live=3600)
-    def _query_total_costs(self, date_range: DateRange, add_attributable_costs_filter):
-        """
-        Internal function to query total costs from AWS Cost Explorer.
-
-        Can query either the total account costs or only the attributable costs
-        based on the add_attributable_costs_filter parameter.
-
-        Args:
-            date_range: DateRange object containing the time period for the query
-            add_attributable_costs_filter: If True, only include attributable costs
-
-        Returns:
-            List of cost entries with 'date', 'cost', and 'name' fields
-        """
-        if add_attributable_costs_filter:
-            name = "attributable"
-            filter = {
+    def query_attributable_costs(self, date_range: DateRange):
+        response = self.query(
+            metrics=[METRICS_UNBLENDED_COST],
+            granularity=GRANULARITY_DAILY,
+            date_range=date_range,
+            filter={
                 "And": [
                     FILTER_USAGE_COSTS,
                     FILTER_ATTRIBUTABLE_COSTS,
                 ]
-            }
-        else:
-            name = "account"
-            filter = FILTER_USAGE_COSTS
-
-        # Use AWS-formatted dates (exclusive end date) for Cost Explorer API
-        from_date, to_date = date_range.aws_range
-
-        response = self.query(
-            metrics=[METRICS_UNBLENDED_COST],
-            granularity=GRANULARITY_DAILY,
-            from_date=from_date,
-            to_date=to_date,
-            filter=filter,
+            },
             group_by=[],
         )
 
         processed_response = [
             {
                 "date": e["TimePeriod"]["Start"],
                 "cost": f"{float(e['Total']['UnblendedCost']['Amount']):.2f}",
-                "name": name,
+                "name": "account",
             }
             for e in response["ResultsByTime"]
         ]
+
         return processed_response
 
     @ttl_lru_cache(seconds_to_live=3600)

tests/data/fixtures/aws-ce/test_query_account_cost-input.json

@@ -0,0 +1,131 @@
+{
+    "GroupDefinitions": [],
+    "ResultsByTime": [
+        {
+            "TimePeriod": {
+                "Start": "2026-06-20",
+                "End": "2026-06-21"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "52.0646447944",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-21",
+                "End": "2026-06-22"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "54.0740726258",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-22",
+                "End": "2026-06-23"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "51.8434995458",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-23",
+                "End": "2026-06-24"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "46.3563414046",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-24",
+                "End": "2026-06-25"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "53.5013421334",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-25",
+                "End": "2026-06-26"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "50.7870556936",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-26",
+                "End": "2026-06-27"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "44.5775389922",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        },
+        {
+            "TimePeriod": {
+                "Start": "2026-06-27",
+                "End": "2026-06-28"
+            },
+            "Total": {
+                "UnblendedCost": {
+                    "Amount": "40.0196681343",
+                    "Unit": "USD"
+                }
+            },
+            "Groups": [],
+            "Estimated": false
+        }
+    ],
+    "DimensionValueAttributes": [],
+    "ResponseMetadata": {
+        "RequestId": "cfc0ac51-43b3-409d-a367-ed9cb4f42740",
+        "HTTPStatusCode": 200,
+        "HTTPHeaders": {
+            "date": "Wed, 29 Jul 2026 22:24:32 GMT",
+            "content-type": "application/x-amz-json-1.1",
+            "content-length": "1295",
+            "connection": "keep-alive",
+            "x-amzn-requestid": "cfc0ac51-43b3-409d-a367-ed9cb4f42740",
+            "cache-control": "no-cache"
+        },
+        "RetryAttempts": 0
+    }
+}

tests/data/fixtures/aws-ce/test_query_account_cost-output.json

@@ -0,0 +1,42 @@
+[
+    {
+        "date": "2026-06-20",
+        "cost": "52.06",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-21",
+        "cost": "54.07",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-22",
+        "cost": "51.84",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-23",
+        "cost": "46.36",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-24",
+        "cost": "53.50",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-25",
+        "cost": "50.79",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-26",
+        "cost": "44.58",
+        "name": "account"
+    },
+    {
+        "date": "2026-06-27",
+        "cost": "40.02",
+        "name": "account"
+    }
+]

tests/test_aws.py

@@ -0,0 +1,30 @@
+import json
+from datetime import datetime
+
+import pytest
+from pytest_httpserver import HTTPServer
+
+from jupyterhub_cost_monitoring.aws import AWSCostExplorer
+from jupyterhub_cost_monitoring.date_utils import DateRange
+
+
+@pytest.fixture
+def aws_date_range() -> DateRange:
+    return DateRange(datetime(2026, 6, 20), datetime(2026, 6, 27))
+
+
+def test_query_account_cost(httpserver: HTTPServer, aws_date_range: DateRange):
+    aws_endpoint_url = f"http://{httpserver.host}:{httpserver.port}/"
+    ce = AWSCostExplorer(
+        aws_client_extra_kwargs={
+            "region_name": "test",
+            "endpoint_url": aws_endpoint_url,
+        }
+    )
+
+    with open("tests/data/fixtures/aws-ce/test_query_account_cost-input.json") as f:
+        httpserver.expect_request("/", method="POST").respond_with_data(f.read())
+
+    account_costs = ce.query_account_costs(aws_date_range)
+    with open("tests/data/fixtures/aws-ce/test_query_account_cost-output.json") as f:
+        assert account_costs == json.load(f)