""" Test cases for o11ydiagnosis/subtasks/slowlogconverter/task.py """ import os import sys import tempfile import shutil import unittest import sqlite3 import time import pickle from concurrent.futures import Future from unittest.mock import patch, MagicMock, Mock, mock_open from datetime import datetime, timezone from io import BytesIO # Add project root to path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../'))) from o11ydiagnosis.subtasks.slowlogconverter.task import ( Task, find_common_ts_prefixes_limit_1h, find_matched_keys, read_keys_from_merged_logs, generate_id, parse_id, deltalake_base_path, log_metrics_periodically, start_metrics_logging, stop_metrics_logging, ) from o11ydiagnosis.subtasks.slowlogconverter.objstore import Local class TestHelperFunctions(unittest.TestCase): """Test helper functions""" def test_find_common_ts_prefixes_limit_1h(self): """Test find_common_ts_prefixes_limit_1h""" begin = 1625097600 # 2021-07-01 10:00:00 UTC end = begin + 3600 # 1 hour later prefixes = find_common_ts_prefixes_limit_1h(begin, end) # Should return prefixes for the hour self.assertGreater(len(prefixes), 0) self.assertTrue(all(isinstance(p, str) for p in prefixes)) def test_find_matched_keys(self): """Test find_matched_keys""" mock_s3 = MagicMock() # Filename must start with 10-digit timestamp for find_matched_keys to work mock_s3.lists.return_value = [ 'logs/slowlog/1625097600/1625097600.log.gz', 'logs/slowlog/1625097600/1625097700.log.gz', 'logs/slowlog/1625098000/1625098000.log.gz', # Out of range ] bucket = 'test-bucket' path = 'logs/slowlog' begin = 1625097600 end = 1625097700 keys = find_matched_keys(mock_s3, bucket, path, begin, end) # Should filter keys by timestamp self.assertGreaterEqual(len(keys), 0) mock_s3.lists.assert_called_once() def test_read_keys_from_merged_logs(self): """Test read_keys_from_merged_logs""" temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.log') temp_file.write("# s3://bucket/key1.log.gz\n") temp_file.write("# s3://bucket/key2.log.gz\n") temp_file.write("normal line\n") temp_file.write("# s3://bucket/key3.log.gz\n") temp_file.close() try: keys = read_keys_from_merged_logs(temp_file.name) self.assertEqual(len(keys), 3) self.assertIn('s3://bucket/key1.log.gz', keys) self.assertIn('s3://bucket/key2.log.gz', keys) self.assertIn('s3://bucket/key3.log.gz', keys) finally: os.unlink(temp_file.name) def test_generate_id(self): """Test generate_id""" dt = datetime(2021, 7, 1, 10, 0, 0, tzinfo=timezone.utc) sequence = 123 id_val = generate_id(dt, sequence) self.assertIsInstance(id_val, int) self.assertGreater(id_val, 0) def test_generate_id_invalid_timestamp(self): """Test generate_id with invalid timestamp""" # Timestamp too large (2^34 = 17179869184, which is around year 2446) # Use a date that will definitely exceed the limit (3000-01-01 is ~32503680000) dt = datetime(3000, 1, 1, 0, 0, 0, tzinfo=timezone.utc) sequence = 0 with self.assertRaises(ValueError): generate_id(dt, sequence) def test_generate_id_invalid_sequence(self): """Test generate_id with invalid sequence""" dt = datetime(2021, 7, 1, 10, 0, 0, tzinfo=timezone.utc) # Sequence too large (>= 2^28) sequence = 1 << 28 with self.assertRaises(ValueError): generate_id(dt, sequence) def test_parse_id(self): """Test parse_id""" dt = datetime(2021, 7, 1, 10, 0, 0, tzinfo=timezone.utc) sequence = 123 id_val = generate_id(dt, sequence) parsed_dt, parsed_sequence = parse_id(id_val) self.assertEqual(parsed_sequence, sequence) # Timestamp should be close (within 1 second) self.assertAlmostEqual(parsed_dt.timestamp(), dt.timestamp(), delta=1) def test_deltalake_base_path_aws(self): """Test deltalake_base_path for AWS""" path = deltalake_base_path("aws", "test-bucket", "test-tenant", "test-cluster") self.assertEqual(path, "s3://test-bucket/deltalake/test-tenant/test-cluster") def test_deltalake_base_path_alicloud(self): """Test deltalake_base_path for AliCloud""" path = deltalake_base_path("alicloud", "test-bucket", "test-tenant", "test-cluster") self.assertEqual(path, "s3://test-bucket/deltalake/test-tenant/test-cluster") def test_deltalake_base_path_azure(self): """Test deltalake_base_path for Azure""" os.environ['AZURE_STORAGE_ACCOUNT'] = 'test-account' try: path = deltalake_base_path("azure", "test-container", "test-tenant", "test-cluster") self.assertIn("abfss://", path) self.assertIn("test-container", path) self.assertIn("test-account", path) finally: if 'AZURE_STORAGE_ACCOUNT' in os.environ: del os.environ['AZURE_STORAGE_ACCOUNT'] def test_deltalake_base_path_gcp(self): """Test deltalake_base_path for GCP""" path = deltalake_base_path("gcp", "test-bucket", "test-tenant", "test-cluster") self.assertEqual(path, "gs://test-bucket/deltalake/test-tenant/test-cluster") def test_deltalake_base_path_local(self): """Test deltalake_base_path for local""" path = deltalake_base_path("local", "test-bucket", "test-tenant", "test-cluster") self.assertEqual(path, "./output/leo") class TestTask(unittest.TestCase): """Test Task class""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.mock_s3 = MagicMock() # Mock metrics self.metrics_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.metrics') self.mock_metrics = self.metrics_patcher.start() self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = None # Mock logger self.logger_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.log') self.mock_log = self.logger_patcher.start() # Mock write functions self.write_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.write_slowlogs_to_deltalake') self.mock_write = self.write_patcher.start() # Mock optimize function self.optimize_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.try_to_optimize_table') self.mock_optimize = self.optimize_patcher.start() # Mock metrics timer self.timer_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task._metrics_timer') self.mock_timer = self.timer_patcher.start() self.mock_timer = None def tearDown(self): """Clean up test fixtures""" self.metrics_patcher.stop() self.logger_patcher.stop() self.write_patcher.stop() self.optimize_patcher.stop() self.timer_patcher.stop() if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def create_sample_slowlog_content(self, num_logs=1): """Create sample slowlog content that can be parsed by StreamParser Based on real slowlog format from o11ydiagnosis/subtasks/slowlogconverter/test_slowlog.py """ lines = [] for i in range(num_logs): timestamp = datetime(2021, 7, 1, 10, 0, i, tzinfo=timezone.utc) # First line format: timestamp;instance;tidb;db-name;slowlog # Time: timestamp # This is the critical format - first line must have "# Time:" on the same line time_str = timestamp.strftime('%Y-%m-%dT%H:%M:%S.%fZ') lines.append(f"2021-07-01T10:00:{i:02d}.000000Z;test-instance;tidb;db-tidb-0;slowlog # Time: {time_str}") lines.append("# Txn_start_ts: 0") lines.append("# User@Host: root[root] @ 127.0.0.1 [127.0.0.1]") lines.append("# Conn_ID: 12345") lines.append("# Query_time: 1.0") lines.append("# Parse_time: 0") lines.append("# Compile_time: 0.001") lines.append("# Rewrite_time: 0.0001") lines.append("# Optimize_time: 0.0005") lines.append("# Wait_TS: 0") lines.append("# Is_internal: false") # Digest must be 64-char hex string digest = 'a' * 64 lines.append(f"# Digest: {digest}") lines.append("# Num_cop_tasks: 0") lines.append("# Prepared: false") lines.append("# Plan_from_cache: false") lines.append("# Plan_from_binding: false") lines.append("# Has_more_results: false") lines.append("# KV_total: 0") lines.append("# PD_total: 0.00001") lines.append("# Backoff_total: 0") lines.append("# Write_sql_response_total: 0.001") lines.append("# Result_rows: 1") lines.append("# Succ: true") lines.append("# IsExplicitTxn: false") # SQL query - must end with semicolon for StreamParser to recognize end of entry lines.append(f"SELECT * FROM test_table_{i} WHERE id = {i};") lines.append("") return '\n'.join(lines).encode('utf-8') def test_task_initialization(self): """Test Task initialization""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix/", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix/", drhouse_local_path=self.temp_dir, ) self.assertEqual(task.tenant, "test-tenant") self.assertEqual(task.cluster, "test-cluster") self.assertEqual(task.o11y_cluster, "test-o11y-cluster") self.assertEqual(task.vendor, "local") self.assertEqual(task.region, "us-west-2") # Prefixes should be stripped of trailing slashes self.assertEqual(task.o11y_prefix, "test/o11y/prefix") self.assertEqual(task.drhouse_prefix, "test/drhouse/prefix") self.assertEqual(task.drhouse_local_path, self.temp_dir) self.assertTrue(os.path.exists(task.temp_dir)) def test_ensure_temp_dir(self): """Test _ensure_temp_dir creates and cleans directory""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) # Create a file in temp_dir test_file = os.path.join(task.temp_dir, "test.txt") with open(test_file, 'w') as f: f.write("test") # Call _ensure_temp_dir again - should clean and recreate task._ensure_temp_dir() # Directory should exist but file should be gone self.assertTrue(os.path.exists(task.temp_dir)) self.assertFalse(os.path.exists(test_file)) @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_no_matched_keys(self, mock_find_keys): """Test run when no keys are matched""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 # 2021-07-01 10:00:00 UTC mock_find_keys.return_value = [] task.run(hour_ts) # Should record metrics with 0 logs self.mock_metrics.record_logs_written.assert_called() mock_find_keys.assert_called_once() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_no_new_keys(self, mock_find_keys): """Test run when all keys are already processed""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Create meta.db with the key already processed metadb_path = os.path.join(self.temp_dir, "meta.db") with sqlite3.connect(metadb_path) as metadb: metadb.isolation_level = None metadb.execute(""" CREATE TABLE IF NOT EXISTS "processed_objects" ( "timestamp" INTEGER NOT NULL, "object_key" TEXT NOT NULL, PRIMARY KEY ("timestamp", "object_key") ); """) metadb.execute( "INSERT INTO processed_objects (timestamp, object_key) VALUES (?, ?)", (hour_ts - hour_ts % 3600, matched_key) ) task.run(hour_ts) # Should record metrics with 0 logs and return early self.mock_metrics.record_logs_written.assert_called() self.mock_s3.download.assert_not_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_parsable_logs(self, mock_find_keys): """Test run with parsable slowlog files using real StreamParser""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Mock S3 download to create a local file with real slowlog content def mock_download(bucket, key, filepath): # Create a file with parsable slowlog content content = self.create_sample_slowlog_content(num_logs=2) with open(filepath, 'wb') as f: f.write(content) return True self.mock_s3.download.side_effect = mock_download # Set LOG_BATCH_SIZE to small value for testing with patch.dict(os.environ, {'LOG_BATCH_SIZE': '1'}): task.run(hour_ts) # Should have downloaded and processed the file self.mock_s3.download.assert_called() # Should have written to deltalake (if logs were parsed successfully) # Note: This depends on StreamParser actually parsing the logs # If parsing fails, write might not be called if self.mock_write.called: self.mock_write.assert_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_real_slowlog_format(self, mock_find_keys): """Test run with real slowlog format that can be parsed by StreamParser""" # Use the exact format from test_slowlog.py that we know works real_slowlog_content = """2022-08-31T07:12:54.233489375Z;test-instance;tidb;db-tidb-0;slowlog # Time: 2022-08-31T07:12:54.096687178Z # Txn_start_ts: 0 # User@Host: root[root] @ 127.0.0.1 [127.0.0.1] # Conn_ID: 12345 # Query_time: 1.0 # Parse_time: 0 # Compile_time: 0.001 # Rewrite_time: 0.0001 # Optimize_time: 0.0005 # Wait_TS: 0 # Is_internal: false # Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb # Num_cop_tasks: 0 # Prepared: false # Plan_from_cache: false # Plan_from_binding: false # Has_more_results: false # KV_total: 0 # PD_total: 0.00001 # Backoff_total: 0 # Write_sql_response_total: 0.001 # Result_rows: 1 # Succ: true # IsExplicitTxn: false SELECT * FROM test_table WHERE id = 1;""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Mock S3 download with real format def mock_download(bucket, key, filepath): with open(filepath, 'wb') as f: f.write(real_slowlog_content.encode('utf-8')) return True self.mock_s3.download.side_effect = mock_download # Set LOG_BATCH_SIZE to small value to trigger batch writes with patch.dict(os.environ, {'LOG_BATCH_SIZE': '1'}): task.run(hour_ts) # Should have downloaded and processed the file self.mock_s3.download.assert_called() # With real format, StreamParser should successfully parse and write should be called self.mock_write.assert_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_optimize_sql_statement(self, mock_find_keys): """Test run with SQL statement optimization""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 mock_find_keys.return_value = [] # Mock existed to return True for SQL statement path sqlstatementpath = f"deltalake/{task.tenant}/{task.o11y_cluster}/persisted_statements_summary/_delta_log/_last_checkpoint" self.mock_s3.existed.return_value = True # Mock last_optimize_timestamp to be None (should trigger optimize) self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = None def _submit_run_inline(fn, *args, **kwargs): fn() fut = Future() fut.set_result(None) return fut # Mock time.time to return a fixed value; run optimize inline (production uses a thread pool) with patch('o11ydiagnosis.subtasks.slowlogconverter.task.OPTIMIZE_EXECUTOR') as mock_pool: mock_pool.submit.side_effect = _submit_run_inline with patch('o11ydiagnosis.subtasks.slowlogconverter.task.time.time', return_value=1625097600): task.run(hour_ts) # Should check if SQL statement path exists (also checks topsql paths afterward) self.mock_s3.existed.assert_any_call(task.o11y_bucket, sqlstatementpath) # Should call optimize if path exists and timestamp is old self.mock_optimize.assert_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_optimize_sql_statement_not_existed(self, mock_find_keys): """Test run when SQL statement path does not exist""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 mock_find_keys.return_value = [] # Mock existed to return False self.mock_s3.existed.return_value = False self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = None task.run(hour_ts) # Should check if SQL statement path exists self.mock_s3.existed.assert_called() # Should not call optimize if path doesn't exist self.mock_optimize.assert_not_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_multiple_batches(self, mock_find_keys): """Test run with multiple file batches (FILE_BATCH_SIZE = 100)""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 # Create more than FILE_BATCH_SIZE (100) keys to test batching matched_keys = [f'logs/slowlog/1625097600/test-{i}.log.gz' for i in range(150)] mock_find_keys.return_value = matched_keys # Mock S3 download def mock_download(bucket, key, filepath): content = self.create_sample_slowlog_content(num_logs=1) with open(filepath, 'wb') as f: f.write(content) return True self.mock_s3.download.side_effect = mock_download # Set LOG_BATCH_SIZE to small value with patch.dict(os.environ, {'LOG_BATCH_SIZE': '1'}): task.run(hour_ts) # Should have downloaded files (at least some) self.assertGreater(self.mock_s3.download.call_count, 0) @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_zero_logs_after_parsing(self, mock_find_keys): """Test run when parsing results in zero logs (invalid content)""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Mock S3 download to create an empty or invalid file def mock_download(bucket, key, filepath): # Create an empty file or file with invalid content with open(filepath, 'wb') as f: f.write(b"invalid content\n") return True self.mock_s3.download.side_effect = mock_download task.run(hour_ts) # Should have downloaded the file self.mock_s3.download.assert_called() # Should record metrics with 0 logs self.mock_metrics.record_logs_written.assert_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_expired_records_cleanup(self, mock_find_keys): """Test run cleans up expired records from processed_objects table""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 mock_find_keys.return_value = [] # Create meta.db with old records metadb_path = os.path.join(self.temp_dir, "meta.db") with sqlite3.connect(metadb_path) as metadb: metadb.isolation_level = None metadb.execute(""" CREATE TABLE IF NOT EXISTS "processed_objects" ( "timestamp" INTEGER NOT NULL, "object_key" TEXT NOT NULL, PRIMARY KEY ("timestamp", "object_key") ); """) # Insert old record (more than 24 hours ago) # begin = hour_ts - hour_ts % 3600, so we need to ensure old_timestamp is before begin - 24*3600 begin = hour_ts - hour_ts % 3600 old_timestamp = begin - 25 * 3600 # More than 24 hours before begin metadb.execute( "INSERT INTO processed_objects (timestamp, object_key) VALUES (?, ?)", (old_timestamp, 'old-key.log.gz') ) task.run(hour_ts) # for there is no new key, the old record should not be deleted # The code deletes records where timestamp < begin - 24 * 3600 begin = hour_ts - hour_ts % 3600 expired_ts = begin - 24 * 3600 with sqlite3.connect(metadb_path) as metadb: cursor = metadb.cursor() cursor.execute("SELECT COUNT(*) FROM processed_objects WHERE timestamp < ?", (expired_ts,)) count = cursor.fetchone()[0] self.assertEqual(count, 1) @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_batch_size_100_in_thread_task(self, mock_find_keys): """Test run with batch size 100 in _thread_task to trigger batch write logic""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Create content with enough logs to trigger batch write (batch size 100 in _thread_task) def mock_download(bucket, key, filepath): # Create multiple slowlog entries (more than 100 to trigger batch write in _thread_task) content_parts = [] for i in range(150): # More than batch size 100 content_parts.append(self.create_sample_slowlog_content(num_logs=1)) content = b'\n'.join(content_parts) with open(filepath, 'wb') as f: f.write(content) return True self.mock_s3.download.side_effect = mock_download # Set LOG_BATCH_SIZE to trigger multiple batch writes with patch.dict(os.environ, {'LOG_BATCH_SIZE': '50'}): task.run(hour_ts) # Should have downloaded and processed self.mock_s3.download.assert_called() # Should have written multiple batches if self.mock_write.called: self.assertGreaterEqual(self.mock_write.call_count, 1) @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_parse_errors(self, mock_find_keys): """Test run when parsing encounters errors (ValueError)""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 matched_key = 'logs/slowlog/1625097600/1625097600.log.gz' mock_find_keys.return_value = [matched_key] # Create file with some valid and some invalid content def mock_download(bucket, key, filepath): # Mix valid and invalid content valid_content = self.create_sample_slowlog_content(num_logs=1) invalid_content = b"invalid log entry\ncorrupted data\n" with open(filepath, 'wb') as f: f.write(valid_content) f.write(invalid_content) return True self.mock_s3.download.side_effect = mock_download with patch.dict(os.environ, {'LOG_BATCH_SIZE': '1'}): task.run(hour_ts) # Should have downloaded the file self.mock_s3.download.assert_called() # Should handle parse errors gracefully self.mock_metrics.record_logs_written.assert_called() @patch('o11ydiagnosis.subtasks.slowlogconverter.task.find_matched_keys') def test_run_with_optimize_triggered_by_time(self, mock_find_keys): """Test run when optimize is triggered (time since last optimize > 3600)""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) hour_ts = 1625097600 mock_find_keys.return_value = [] # Mock existed to return True self.mock_s3.existed.return_value = True # Mock last_optimize_timestamp to be old (more than 3600 seconds ago) old_timestamp = 1625097600 - 7200 # 2 hours ago self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = old_timestamp # Mock time.time to return current time with patch('o11ydiagnosis.subtasks.slowlogconverter.task.time.time', return_value=1625097600): task.run(hour_ts) # Should call optimize since timestamp is old self.mock_optimize.assert_called() def test_need_optimize_future_timestamp_forced_true(self): """Future last_optimize_timestamp should not block optimize scheduling.""" task = Task( s3=self.mock_s3, tenant="test-tenant", cluster="test-cluster", o11y_cluster="test-o11y-cluster", vendor="local", region="us-west-2", o11y_bucket="test-o11y-bucket", o11y_prefix="test/o11y/prefix", drhouse_bucket="test-drhouse-bucket", drhouse_prefix="test/drhouse/prefix", drhouse_local_path=self.temp_dir, ) now = 1625097600 self.mock_metrics.last_optimize_timestamp.labels.return_value._value.get.return_value = now + 3600 with patch('o11ydiagnosis.subtasks.slowlogconverter.task.time.time', return_value=now): self.assertTrue(task._need_optimize()) self.mock_log.warning.assert_called() class TestMetricsLogging(unittest.TestCase): """Test metrics logging functions""" def setUp(self): """Set up test fixtures""" # Mock metrics self.metrics_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.metrics') self.mock_metrics = self.metrics_patcher.start() self.mock_metrics._cluster_hours = {} self.mock_metrics.get_oldest_processed_hour.return_value = None self.mock_metrics.last_update._metrics = {} self.mock_metrics.logs_written_1m._metrics = {} # Mock logger self.logger_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.log') self.mock_log = self.logger_patcher.start() # Mock pympler self.pympler_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.muppy') self.mock_muppy = self.pympler_patcher.start() self.mock_muppy.get_objects.return_value = [] self.summary_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.summary') self.mock_summary = self.summary_patcher.start() self.mock_summary.summarize.return_value = [] # Mock Timer self.timer_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.task.Timer') self.mock_timer_class = self.timer_patcher.start() self.mock_timer = MagicMock() self.mock_timer_class.return_value = self.mock_timer def tearDown(self): """Clean up test fixtures""" self.metrics_patcher.stop() self.logger_patcher.stop() self.pympler_patcher.stop() self.summary_patcher.stop() self.timer_patcher.stop() def test_log_metrics_periodically(self): """Test log_metrics_periodically function""" # Mock oldest processed hour self.mock_metrics.get_oldest_processed_hour.return_value = ("test-cluster", 1625097600) log_metrics_periodically() # Should log metrics self.mock_log.info.assert_called() # Should schedule next run self.mock_timer.start.assert_called() def test_start_metrics_logging(self): """Test start_metrics_logging function""" # Mock existing timer with patch('o11ydiagnosis.subtasks.slowlogconverter.task._metrics_timer', None): start_metrics_logging() self.mock_timer_class.assert_called() def test_stop_metrics_logging(self): """Test stop_metrics_logging function""" # Mock existing timer with patch('o11ydiagnosis.subtasks.slowlogconverter.task._metrics_timer', self.mock_timer): stop_metrics_logging() self.mock_timer.cancel.assert_called_once() if __name__ == '__main__': unittest.main()