import os import shutil import tempfile import unittest from unittest.mock import patch, MagicMock, ANY from datetime import datetime, timezone import pyarrow as pa import pandas as pd try: # Try to import from the installed deltalake package import sys # Remove local deltalake from path temporarily if 'deltalake' in sys.modules: del sys.modules['deltalake'] from deltalake import DeltaTable, write_deltalake DELTALAKE_AVAILABLE = True except ImportError: DELTALAKE_AVAILABLE = False DeltaTable = None write_deltalake = None from o11ydiagnosis.subtasks.slowlogconverter.slowlog import SlowLog from o11ydiagnosis.subtasks.slowlogconverter.write_deltalake import ( write_slowlogs_to_deltalake, write_slowlogs_to_deltalake_multitables, slowlogs_schema, plans_schema, queries_schema, ) class TestWriteDeltaLake(unittest.TestCase): @unittest.skipUnless(DELTALAKE_AVAILABLE, "deltalake package not available") def setUp(self): # Create a temporary directory for test Delta tables self.temp_dir = tempfile.mkdtemp() self.base_path = os.path.join(self.temp_dir, "deltalake") # Create a sample SlowLog object for testing self.sample_log = SlowLog() self.sample_log.id = 12345 self.sample_log.time = 1625097600 # 2021-06-30 16:00:00 UTC self.sample_log.digest = "test_digest_123" self.sample_log.instance = "test-instance-1" self.sample_log.query = "SELECT * FROM test_table" self.sample_log.binary_plan = "test_binary_plan" self.sample_log.plan = "test_plan" self.sample_log.plan_digest = "test_plan_digest_123" self.sample_log.txn_start_ts = 1234567890123456789 self.sample_log.conn_id = "12345" self.sample_log.exec_retry_count = 0 self.sample_log.preproc_subqueries = 0 self.sample_log.request_count = 1 self.sample_log.total_keys = 10 self.sample_log.process_keys = 5 self.sample_log.rocksdb_delete_skipped_count = 0 self.sample_log.rocksdb_key_skipped_count = 0 self.sample_log.rocksdb_block_cache_hit_count = 5 self.sample_log.rocksdb_block_read_count = 5 self.sample_log.rocksdb_block_read_byte = 1024 self.sample_log.db = "test_db" # Add required fields self.sample_log.index_names = "" self.sample_log.is_internal = False self.sample_log.stats = "" self.sample_log.cop_proc_avg = 0.0 self.sample_log.cop_proc_p90 = 0.0 self.sample_log.cop_proc_max = 0.0 self.sample_log.cop_proc_addr = "" self.sample_log.cop_wait_avg = 0.0 self.sample_log.cop_wait_p90 = 0.0 self.sample_log.cop_wait_max = 0.0 self.sample_log.cop_wait_addr = "" self.sample_log.mem_max = 0 self.sample_log.disk_max = 0 self.sample_log.kv_total = 0.0 # Mock the parse_binary_plan function self.parse_binary_plan_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.parse_binary_plan') self.mock_parse_binary_plan = self.parse_binary_plan_patcher.start() self.mock_parse_binary_plan.side_effect = lambda x: f"decoded_{x}" if x else "" # Mock the logger self.logger_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.logger') self.mock_logger = self.logger_patcher.start() # Mock PyArrow Table creation self.pyarrow_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.pa.Table') self.mock_pyarrow_table = self.pyarrow_patcher.start() self.mock_pyarrow_table.from_pandas.return_value = MagicMock() # Mock the DeltaTable and write_deltalake self.deltalake_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.DeltaTable') self.mock_delta_table = self.deltalake_patcher.start() self.mock_delta_table.return_value.merge.return_value.when_matched_update_all.return_value.when_not_matched_insert_all.return_value.execute.return_value = None # Mock write_deltalake self.write_deltalake_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.write_deltalake') self.mock_write_deltalake = self.write_deltalake_patcher.start() def tearDown(self): """Tear down test fixtures.""" # Clean up the temporary directory if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) # Stop all patches self.parse_binary_plan_patcher.stop() self.logger_patcher.stop() self.deltalake_patcher.stop() self.write_deltalake_patcher.stop() self.pyarrow_patcher.stop() def test_write_slowlogs_to_deltalake(self): """Test writing slow logs to a single Delta table""" # Prepare test data test_logs = [self.sample_log] # Call the function write_slowlogs_to_deltalake( added_logs=test_logs, region="us-west-2", cluster="test-cluster", base_path=self.base_path, optimize=False, # Skip optimization for faster tests merge_write=True ) # Verify the Delta table was written to expected_path = f"{self.base_path}/slowlogs" # Verify parse_binary_plan was called with the correct argument self.mock_parse_binary_plan.assert_called_once_with(self.sample_log.binary_plan) # Merge mode is disabled; always use append + schema evolution. self.mock_delta_table.return_value.merge.assert_not_called() self.mock_write_deltalake.assert_called_once() call_args = self.mock_write_deltalake.call_args[1] self.assertEqual(call_args["mode"], "append") self.assertEqual(call_args["schema_mode"], "merge") # Verify the logger was called with success message (allow partial match) calls = [str(call) for call in self.mock_logger.info.call_args_list] self.assertTrue(any("[test-cluster]" in str(call) and "write_slowlogs_data_end, cost" in str(call) for call in calls), f"Expected log message not found. Actual calls: {calls}") def test_write_slowlogs_to_deltalake_multitables(self): """Test writing slow logs to multiple Delta tables (slowlogs, queries, plans)""" test_logs = [self.sample_log] write_slowlogs_to_deltalake_multitables( added_logs=test_logs, region="us-west-2", cluster="test-cluster", base_path=self.base_path, optimize=False, merge_write=True ) # Verify all three tables were written to expected_paths = [ f"{self.base_path}/slowlogs", f"{self.base_path}/queries", f"{self.base_path}/plans" ] # Merge mode is disabled; append is used for each table. self.mock_delta_table.return_value.merge.assert_not_called() self.assertEqual(self.mock_write_deltalake.call_count, 3) # Verify the logger was called with success messages (allow partial match) calls = [str(call) for call in self.mock_logger.info.call_args_list] self.assertTrue(any("[test-cluster]" in str(call) and "write_slowlogs_data_end, cost" in str(call) for call in calls), f"Expected log message not found. Actual calls: {calls}") self.assertTrue(any("[test-cluster]" in str(call) and "write_queries_data_end, cost" in str(call) for call in calls), f"Expected log message not found. Actual calls: {calls}") self.assertTrue(any("[test-cluster]" in str(call) and "write_plans_data_end, cost" in str(call) for call in calls), f"Expected log message not found. Actual calls: {calls}") def test_write_slowlogs_to_deltalake_with_partition_date(self): """Test writing with a specific partition date""" test_logs = [self.sample_log] partition_date = "20210701" write_slowlogs_to_deltalake( added_logs=test_logs, region="us-west-2", cluster="test-cluster", base_path=self.base_path, optimize=False, partition_date=partition_date, merge_write=True ) # Merge mode is disabled; verify partition columns are passed. self.mock_delta_table.return_value.merge.assert_not_called() self.mock_write_deltalake.assert_called_once() call_args = self.mock_write_deltalake.call_args[1] self.assertEqual(call_args["mode"], "append") self.assertEqual(call_args["schema_mode"], "merge") self.assertEqual(call_args["partition_by"], ["date"]) def test_write_slowlogs_to_deltalake_empty_logs(self): """Test writing an empty list of logs (should not fail)""" write_slowlogs_to_deltalake( added_logs=[], region="us-west-2", cluster="test-cluster", base_path=self.base_path, optimize=False, merge_write=True ) # Empty logs should not trigger any write. self.mock_write_deltalake.assert_not_called() def test_write_slowlogs_to_deltalake_optimize(self): """Test writing with optimization enabled""" # Mock the optimize functions with patch('o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.try_to_optimize_table') as mock_optimize: test_logs = [self.sample_log] write_slowlogs_to_deltalake( added_logs=test_logs, region="us-west-2", cluster="test-cluster", base_path=self.base_path, optimize=True, merge_write=True ) # Verify optimization was attempted mock_optimize.assert_called_once() if __name__ == '__main__': unittest.main()