From d02f18660ea1f58550cd238cbc87b16cd28dc895 Mon Sep 17 00:00:00 2001
From: TeachMeTW <robin@robinttw.com>
Date: Tue, 15 Oct 2024 10:11:38 -0700
Subject: [PATCH] Used get entry at ts

---
 .../storage/timeseries/builtin_timeseries.py  | 13 ++++----
 .../tests/funcTests/TestFunctionTiming.py     | 32 +++++--------------
 2 files changed, 15 insertions(+), 30 deletions(-)

diff --git a/emission/storage/timeseries/builtin_timeseries.py b/emission/storage/timeseries/builtin_timeseries.py
index 037951e9b..9bd8669d4 100644
--- a/emission/storage/timeseries/builtin_timeseries.py
+++ b/emission/storage/timeseries/builtin_timeseries.py
@@ -170,12 +170,13 @@ def _get_query(self, key_list = None, time_query = None, geo_query = None,
                 eq_keys = set(extra_query.keys())
                 curr_keys = set(ret_query.keys())
                 overlap_keys = eq_keys.intersection(curr_keys)
-                if overlap_keys:
-                    logging.info(
-                        f"eq_keys = {eq_keys}, curr_keys = {curr_keys}, overlap_keys = {overlap_keys}"
-                    )
-                    raise AttributeError(f"extra query would overwrite keys {list(overlap_keys)}")
-                ret_query.update(extra_query)
+                if len(overlap_keys) != 0:
+                    logging.info("eq_keys = %s, curr_keys = %s, overlap_keys = %s" %
+                                 (eq_keys, curr_keys, overlap_keys))
+                    raise AttributeError("extra query would overwrite keys %s" %
+                                         list(overlap_keys))
+                else:
+                    ret_query.update(extra_query)
 
         return ret_query
 
diff --git a/emission/tests/funcTests/TestFunctionTiming.py b/emission/tests/funcTests/TestFunctionTiming.py
index e5c16eb88..534eafefe 100644
--- a/emission/tests/funcTests/TestFunctionTiming.py
+++ b/emission/tests/funcTests/TestFunctionTiming.py
@@ -70,19 +70,11 @@ def execute_and_time_function(func: t.Callable[[], bool]):
         # Verification: Adjusted Query to Match Document Structure
         timeseries_db = esta.TimeSeries.get_time_series(None)
 
-        # Define the time range and additional filters
-        time_query = {
-            "data.ts": {"$gte": timestamp, "$lte": timestamp},
-            "data.name": function_name,
-            "data.reading": {"$gte": elapsed_ms, "$lte": elapsed_ms}
-        }
-
-        # Retrieve the first matching entry using the get_first_entry method
-        stored_document = timeseries_db.get_first_entry(
+        # Retrieve the document
+        stored_document = timeseries_db.get_entry_at_ts(
             key="stats/dashboard_time",
-            field="data.ts",
-            sort_order=pymongo.DESCENDING,
-            time_query=time_query
+            ts_key="data.ts",
+            ts=timestamp
         )
 
         if stored_document:
@@ -120,19 +112,11 @@ def execute_and_time_function(func: t.Callable[[], bool]):
         # Initialize the TimeSeries database connection
         timeseries_db = esta.TimeSeries.get_time_series(None)
 
-        # Define the time range and additional filters
-        time_query = {
-            "data.ts": {"$gte": timestamp, "$lte": timestamp},
-            "data.name": function_name,
-            "data.reading": {"$gte": elapsed_ms, "$lte": elapsed_ms}
-        }
-
-        # Retrieve the first matching entry using the get_first_entry method
-        stored_error = timeseries_db.get_first_entry(
+        # Retrieve the document
+        stored_error = timeseries_db.get_entry_at_ts(
             key="stats/dashboard_error",
-            field="data.ts",
-            sort_order=pymongo.DESCENDING,
-            time_query=time_query
+            ts_key="data.ts",
+            ts=timestamp
         )