From f57c5d681d3c3305d39850847e5de116e009cb33 Mon Sep 17 00:00:00 2001
From: Rambatino <mark.tint.ramotowski@gmail.com>
Date: Wed, 14 Feb 2024 14:46:10 +0000
Subject: [PATCH] (tests): fixed the tests

---
 axiom/client.py       | 10 +++++-----
 axiom/query/result.py | 14 +++++++-------
 tests/test_client.py  |  2 +-
 tests/test_logger.py  |  4 ++--
 4 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/axiom/client.py b/axiom/client.py
index a121ce4..f71e76a 100644
--- a/axiom/client.py
+++ b/axiom/client.py
@@ -253,7 +253,7 @@ def query_legacy(
         self.logger.debug(f"query result: {result}")
         query_id = res.headers.get("X-Axiom-History-Query-Id")
         self.logger.info(f"received query result with query_id: {query_id}")
-        result.saved_query_id = query_id
+        result.query_id = query_id
         return result
 
     def apl_query(
@@ -276,16 +276,16 @@ def query(
         res = self.session.post(path, data=payload, params=params)
         result = Util.from_dict(
             (
-                TabularQueryResult
-                if opts.format == AplResultFormat.Tabular
-                else LegacyQueryResult
+                LegacyQueryResult
+                if opts is None or opts.format == AplResultFormat.Legacy
+                else TabularQueryResult
             ),
             res.json(),
         )
         self.logger.debug(f"apl query result: {result}")
         query_id = res.headers.get("X-Axiom-History-Query-Id")
         self.logger.info(f"received query result with query_id: {query_id}")
-        result.saved_query_id = query_id
+        result.query_id = query_id
         return result
 
     def df(
diff --git a/axiom/query/result.py b/axiom/query/result.py
index e0b8a2a..5277cd6 100644
--- a/axiom/query/result.py
+++ b/axiom/query/result.py
@@ -145,10 +145,10 @@ class QueryLegacyResult:
     matches: List[Entry]
     # Buckets are the time series buckets.
     buckets: Timeseries
-    # save_query_id is the ID of the query that generated this result when it
+    # query_id is the ID of the query that generated this result when it
     # was saved on the server. This is only set when the query was sent with
     # the `save_as_kind` option specified.
-    save_query_id: Optional[str] = field(default=None)
+    query_id: Optional[str] = field(default=None)
 
 
 @dataclass
@@ -164,10 +164,10 @@ class LegacyQueryResult:
     buckets: Timeseries
     # Dataset names are the datasets that were used in the apl query.
     dataset_names: List[str] = field(default_factory=lambda: [])
-    # save_query_id is the ID of the apl query that generated this result when it
+    # query_id is the ID of the apl query that generated this result when it
     # was saved on the server. This is only set when the apl query was sent with
     # the `save_as_kind` option specified.
-    save_query_id: Optional[str] = field(default=None)
+    query_id: Optional[str] = field(default=None)
 
 
 @dataclass
@@ -294,14 +294,14 @@ class TabularQueryResult:
         request (Request): The request that generated this result.
         dataset_names (List[str]): The names of datasets included in the result.
         fields_meta_map (Dict[str, List[Any]]): Metadata for the fields in the result.
-        save_query_id (Optional[str]): The ID of the saved query that generated this result, if applicable.
+        query_id (Optional[str]): The ID of the saved query that generated this result, if applicable.
     """
 
     format: str
     status: QueryStatus
     tables: List[Table]
     request: Request
-    # save_query_id is the ID of the query that generated this result when it
+    # query_id is the ID of the query that generated this result when it
     # was saved on the server. This is only set when the query was sent with
     # the `save_as_kind` option specified.
-    save_query_id: Optional[str] = field(default=None)
+    query_id: Optional[str] = field(default=None)
diff --git a/tests/test_client.py b/tests/test_client.py
index 4c287fd..a0f3132 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -155,7 +155,7 @@ def test_step004_query(self):
         )
         qr = self.client.query_legacy(self.dataset_name, q, opts)
 
-        self.assertIsNotNone(qr.save_query_id)
+        self.assertIsNotNone(qr.query_id)
         self.assertEqual(len(qr.matches), len(self.events))
 
     def test_step005_apl_query(self):
diff --git a/tests/test_logger.py b/tests/test_logger.py
index 68eed92..4bba1be 100644
--- a/tests/test_logger.py
+++ b/tests/test_logger.py
@@ -32,14 +32,14 @@ def test_log(self):
 
         # this log shouldn't be ingested yet
         res = client.apl_query(dataset_name)
-        self.assertEqual(0, res.status.rowsExamined)
+        self.assertEqual(0, res.status.rows_examined)
 
         # flush events
         axiom_handler.flush()
 
         # now we should have a log
         res = client.apl_query(dataset_name)
-        self.assertEqual(1, res.status.rowsExamined)
+        self.assertEqual(1, res.status.rows_examined)
 
         # cleanup created dataset
         client.datasets.delete(dataset_name)