Fixing unit tests

DAGWorks-Inc · May 8, 2024 · 1e5a616 · 1e5a616
1 parent 4727182
commit 1e5a616
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/ui/sdk/tests/tracking/test_stats.py b/ui/sdk/tests/tracking/test_stats.py
@@ -1,3 +1,4 @@
+import sys
 from collections import namedtuple
 from typing import NamedTuple
 
@@ -41,15 +42,22 @@ def test_compute_stats_tuple_dataloader():
         "test_node",
         {"hamilton.data_loader": True},
     )
+    # in python 3.11 numbers change
+    if sys.version_info >= (3, 11):
+        index = 132
+        memory = 156
+    else:
+        index = 128
+        memory = 152
     assert actual == {
         "observability_schema_version": "0.0.2",
         "observability_type": "dict",
         "observability_value": {
             "type": "<class 'dict'>",
             "value": {
                 "CONNECTION_INFO": {"URL": "SOME_URL"},
-                "DF_MEMORY_BREAKDOWN": {"Index": 128, "a": 24},
-                "DF_MEMORY_TOTAL": 152,
+                "DF_MEMORY_BREAKDOWN": {"Index": index, "a": 24},
+                "DF_MEMORY_TOTAL": memory,
                 "QUERIED_TABLES": {
                     "table-0": {"catalog": "FOO", "database": "BAR", "name": "TABLE"}
                 },