fix unit tests

ray-project · Jun 27, 2023 · a213926 · a213926
1 parent 77f058b
commit a213926
Showing 1 changed file with 22 additions and 5 deletions.
diff --git a/deltacat/tests/compactor/utils/test_io.py b/deltacat/tests/compactor/utils/test_io.py
@@ -1,9 +1,14 @@
 import unittest
 from unittest import mock
+from deltacat.compute.compactor.model.compaction_session_audit_info import (
+    CompactionSessionAuditInfo,
+)
 from deltacat.tests.test_utils.constants import TEST_DELTA
 
 
 class TestFitInputDeltas(unittest.TestCase):
+    COMPACTION_AUDIT = CompactionSessionAuditInfo()
+
     @classmethod
     def setUpClass(cls):
         cls.module_patcher = mock.patch.dict("sys.modules", {"ray": mock.MagicMock()})
@@ -19,12 +24,18 @@ def test_sanity(self):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, None)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, None
+        )
 
         self.assertIsNotNone(hash_bucket_count)
         self.assertTrue(1, len(delta_list))
         self.assertIsNotNone(high_watermark)
         self.assertFalse(require_multiple_rounds)
+        self.assertIsNotNone(hash_bucket_count, self.COMPACTION_AUDIT.hash_bucket_count)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.input_file_count)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.input_size_bytes)
+        self.assertIsNotNone(self.COMPACTION_AUDIT.total_cluster_memory_bytes)
 
     def test_when_hash_bucket_count_overridden(self):
         from deltacat.compute.compactor.utils import io
@@ -34,7 +45,9 @@ def test_when_hash_bucket_count_overridden(self):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 1, "memory": 20000000}, 20)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 1, "memory": 20000000}, self.COMPACTION_AUDIT, 20
+        )
 
         self.assertEqual(20, hash_bucket_count)
         self.assertEqual(1, len(delta_list))
@@ -49,7 +62,9 @@ def test_when_not_enough_memory_splits_manifest_entries(self):
             hash_bucket_count,
             high_watermark,
             require_multiple_rounds,
-        ) = io.fit_input_deltas([TEST_DELTA], {"CPU": 2, "memory": 10}, 20)
+        ) = io.fit_input_deltas(
+            [TEST_DELTA], {"CPU": 2, "memory": 10}, self.COMPACTION_AUDIT, 20
+        )
 
         self.assertIsNotNone(hash_bucket_count)
         self.assertTrue(2, len(delta_list))
@@ -60,10 +75,12 @@ def test_when_no_input_deltas(self):
         from deltacat.compute.compactor.utils import io
 
         with self.assertRaises(AssertionError):
-            io.fit_input_deltas([], {"CPU": 100, "memory": 20000.0}, None)
+            io.fit_input_deltas(
+                [], {"CPU": 100, "memory": 20000.0}, self.COMPACTION_AUDIT, None
+            )
 
     def test_when_cpu_resources_is_not_passed(self):
         from deltacat.compute.compactor.utils import io
 
         with self.assertRaises(KeyError):
-            io.fit_input_deltas([], {}, None)
+            io.fit_input_deltas([], {}, self.COMPACTION_AUDIT, None)