""" Test cases for o11ydiagnosis/subtasks/slowlogcompactor/task.py """ import os import sys import tempfile import shutil import unittest import sqlite3 from unittest.mock import patch, MagicMock from datetime import datetime # Add project root to path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../'))) from o11ydiagnosis.subtasks.slowlogcompactor.task import Task from o11ydiagnosis.subtasks.slowlogconverter.db import StatsDB from o11ydiagnosis.subtasks.slowlogconverter.slowstats import SlowStats, merge_stats from o11ydiagnosis.subtasks.slowlogconverter.objstore import Local class TestTask(unittest.TestCase): """Test Task class""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.local_path = os.path.join(self.temp_dir, "local") os.makedirs(self.local_path, exist_ok=True) self.mock_s3 = MagicMock() self.cluster = "test-cluster" self.bucket = "test-bucket" self.prefix = "test-prefix" self.task = Task( self.mock_s3, self.cluster, self.bucket, self.prefix, self.local_path ) def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_task_initialization(self): """Test Task initialization""" self.assertEqual(self.task.s3, self.mock_s3) self.assertEqual(self.task.cluster, self.cluster) self.assertEqual(self.task.drhouse_bucket, self.bucket) self.assertEqual(self.task.drhouse_prefix, self.prefix) self.assertEqual(self.task.drhouse_local_path, self.local_path) self.assertEqual(self.task.existed_statsdb_paths, {}) self.assertEqual(self.task.dirty_merged_statsdb_paths, {}) def test_task_initialization_strips_trailing_slashes(self): """Test that trailing slashes are stripped""" task = Task( self.mock_s3, self.cluster, self.bucket, self.prefix + "/", self.local_path + "/" ) self.assertEqual(task.drhouse_prefix, self.prefix) self.assertEqual(task.drhouse_local_path, self.local_path) def create_test_statsdb(self, filepath, keys=None, meta=None, stats=None, stats_by_db=None, count_per_minute=None): """Helper to create a test statsdb file""" os.makedirs(os.path.dirname(filepath), exist_ok=True) with sqlite3.connect(filepath) as db: StatsDB.try_init_tables(db) if keys: StatsDB.save_keys(db, keys) if meta: StatsDB.save_meta(db, meta) if stats: StatsDB.save_stats(db, stats) if stats_by_db: StatsDB.save_stats_by_db(db, stats_by_db) if count_per_minute: StatsDB.save_count_per_minute(db, count_per_minute) db.commit() @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_prepare_downloads_files(self, mock_log): """Test prepare downloads statsdb files""" begin_hour_ts = 1609459200 # 2021-01-01 00:00:00 UTC n_hours = 2 self.mock_s3.download.return_value = True self.task.prepare(begin_hour_ts, n_hours) # Should attempt to download n_hours files self.assertEqual(self.mock_s3.download.call_count, n_hours) # Check that paths are tracked self.assertEqual(len(self.task.existed_statsdb_paths), n_hours) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_prepare_uses_existing_files(self, mock_log): """Test prepare uses existing local files""" begin_hour_ts = 1609459200 n_hours = 2 # Create a local file utc_hour = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") statsdb_local_path = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour}.db" ) self.create_test_statsdb(statsdb_local_path, keys=["key1"]) self.task.prepare(begin_hour_ts, n_hours) # Should not download if file exists # But still should be tracked self.assertIn(statsdb_local_path, self.task.existed_statsdb_paths) # Should only download the missing file self.assertEqual(self.mock_s3.download.call_count, 1) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_prepare_asserts_n_hours(self, mock_log): """Test prepare asserts n_hours > 1""" begin_hour_ts = 1609459200 with self.assertRaises(AssertionError): self.task.prepare(begin_hour_ts, 1) with self.assertRaises(AssertionError): self.task.prepare(begin_hour_ts, 0) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_prepare_asserts_timestamp(self, mock_log): """Test prepare asserts timestamp is divisible by n_hours * 3600""" begin_hour_ts = 1609459200 n_hours = 2 # Invalid timestamp (not divisible by n_hours * 3600) invalid_ts = begin_hour_ts + 3600 with self.assertRaises(AssertionError): self.task.prepare(invalid_ts, n_hours) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_prepare_handles_missing_s3_files(self, mock_log): """Test prepare handles missing S3 files gracefully""" begin_hour_ts = 1609459200 n_hours = 2 # Mock download to return False (file doesn't exist) self.mock_s3.download.return_value = False # Should not raise exception self.task.prepare(begin_hour_ts, n_hours) # Should still attempt downloads self.assertEqual(self.mock_s3.download.call_count, n_hours) # But files won't be tracked if download fails self.assertEqual(len(self.task.existed_statsdb_paths), 0) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_merges_multiple_databases(self, mock_log): """Test run merges data from multiple databases""" begin_hour_ts = 1609459200 n_hours = 2 # Create two test databases utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") utc_hour2 = datetime.utcfromtimestamp(begin_hour_ts + 3600).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) statsdb_path2 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour2}.db" ) # Create first database with some data stats1 = SlowStats("digest1") stats1.count = 10 stats1.query_time_sum = 100.0 stats1.query_time_max = 20.0 self.create_test_statsdb( statsdb_path1, keys=["key1", "key2"], meta={"digest1": ("SELECT * FROM test", False)}, stats={"digest1": stats1}, stats_by_db={"db1": {"digest1": stats1}}, count_per_minute={"digest1": {1609459200: 5, 1609459260: 5}} ) # Create second database with overlapping and new data stats2 = SlowStats("digest1") stats2.count = 5 stats2.query_time_sum = 50.0 stats2.query_time_max = 15.0 stats3 = SlowStats("digest2") stats3.count = 3 stats3.query_time_sum = 30.0 self.create_test_statsdb( statsdb_path2, keys=["key2", "key3"], meta={"digest1": ("SELECT * FROM test", False), "digest2": ("INSERT INTO test", False)}, stats={"digest1": stats2, "digest2": stats3}, stats_by_db={"db1": {"digest1": stats2}}, count_per_minute={"digest1": {1609459200: 3, 1609459320: 2}} ) # Track the databases self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.existed_statsdb_paths[statsdb_path2] = "s3-key2" self.task.run(begin_hour_ts, n_hours) # Check that merged database was created merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) self.assertTrue(os.path.exists(merged_path)) # Verify merged data with sqlite3.connect(merged_path) as db: merged_keys = StatsDB.query_keys(db) merged_meta = StatsDB.query_meta(db) merged_stats = StatsDB.query_stats(db) merged_stats_by_db = StatsDB.query_stats_by_db(db) merged_count_per_minute = StatsDB.query_count_per_minute(db) # Keys should be union self.assertEqual(set(merged_keys), {"key1", "key2", "key3"}) # Meta should be merged self.assertEqual(len(merged_meta), 2) self.assertIn("digest1", merged_meta) self.assertIn("digest2", merged_meta) # Stats should be merged (digest1 should have combined counts) self.assertIn("digest1", merged_stats) self.assertEqual(merged_stats["digest1"].count, 15) # 10 + 5 self.assertEqual(merged_stats["digest1"].query_time_sum, 150.0) # 100 + 50 self.assertEqual(merged_stats["digest1"].query_time_max, 20.0) # max(20, 15) # Check that merged database is tracked self.assertIn(merged_path, self.task.dirty_merged_statsdb_paths) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_no_data_skips_creation(self, mock_log): """Test run skips database creation when no data""" begin_hour_ts = 1609459200 n_hours = 2 # Create empty databases utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) self.create_test_statsdb(statsdb_path1, keys=[]) # Empty keys self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.run(begin_hour_ts, n_hours) # Should not create merged database merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) self.assertFalse(os.path.exists(merged_path)) self.assertEqual(len(self.task.dirty_merged_statsdb_paths), 0) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_merges_stats_correctly(self, mock_log): """Test run merges stats correctly using merge_stats""" begin_hour_ts = 1609459200 n_hours = 2 utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") utc_hour2 = datetime.utcfromtimestamp(begin_hour_ts + 3600).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) statsdb_path2 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour2}.db" ) # Create stats with different max values stats1 = SlowStats("digest1") stats1.count = 10 stats1.query_time_max = 20.0 stats1.mem_max_max = 100.0 stats1.disk_max_max = 200.0 stats2 = SlowStats("digest1") stats2.count = 5 stats2.query_time_max = 25.0 # Higher max stats2.mem_max_max = 150.0 # Higher max stats2.disk_max_max = 180.0 # Lower max self.create_test_statsdb(statsdb_path1, keys=["key1"], stats={"digest1": stats1}) self.create_test_statsdb(statsdb_path2, keys=["key2"], stats={"digest1": stats2}) self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.existed_statsdb_paths[statsdb_path2] = "s3-key2" self.task.run(begin_hour_ts, n_hours) # Verify merged stats merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) with sqlite3.connect(merged_path) as db: merged_stats = StatsDB.query_stats(db) self.assertIn("digest1", merged_stats) self.assertEqual(merged_stats["digest1"].count, 15) # 10 + 5 self.assertEqual(merged_stats["digest1"].query_time_max, 25.0) # max(20, 25) self.assertEqual(merged_stats["digest1"].mem_max_max, 150.0) # max(100, 150) self.assertEqual(merged_stats["digest1"].disk_max_max, 200.0) # max(200, 180) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_merges_count_per_minute(self, mock_log): """Test run merges count_per_minute correctly""" begin_hour_ts = 1609459200 n_hours = 2 utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") utc_hour2 = datetime.utcfromtimestamp(begin_hour_ts + 3600).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) statsdb_path2 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour2}.db" ) # Create databases with overlapping timestamps count1 = { "digest1": {1609459200: 5, 1609459260: 3} } count2 = { "digest1": {1609459200: 2, 1609459320: 4} # Overlapping timestamp } self.create_test_statsdb(statsdb_path1, keys=["key1"], count_per_minute=count1) self.create_test_statsdb(statsdb_path2, keys=["key2"], count_per_minute=count2) self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.existed_statsdb_paths[statsdb_path2] = "s3-key2" self.task.run(begin_hour_ts, n_hours) # Verify merged count_per_minute merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) with sqlite3.connect(merged_path) as db: merged_count = StatsDB.query_count_per_minute(db) self.assertIn("digest1", merged_count) # Overlapping timestamp should be summed self.assertEqual(merged_count["digest1"][1609459200], 7) # 5 + 2 self.assertEqual(merged_count["digest1"][1609459260], 3) self.assertEqual(merged_count["digest1"][1609459320], 4) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_merges_stats_by_db(self, mock_log): """Test run merges stats_by_db correctly""" begin_hour_ts = 1609459200 n_hours = 2 utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") utc_hour2 = datetime.utcfromtimestamp(begin_hour_ts + 3600).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) statsdb_path2 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour2}.db" ) stats1 = SlowStats("digest1") stats1.count = 10 stats2 = SlowStats("digest1") stats2.count = 5 stats3 = SlowStats("digest2") stats3.count = 3 self.create_test_statsdb( statsdb_path1, keys=["key1"], stats_by_db={"db1": {"digest1": stats1}} ) self.create_test_statsdb( statsdb_path2, keys=["key2"], stats_by_db={"db1": {"digest1": stats2}, "db2": {"digest2": stats3}} ) self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.existed_statsdb_paths[statsdb_path2] = "s3-key2" self.task.run(begin_hour_ts, n_hours) # Verify merged stats_by_db merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) with sqlite3.connect(merged_path) as db: merged_stats_by_db = StatsDB.query_stats_by_db(db) self.assertIn("db1", merged_stats_by_db) self.assertIn("db2", merged_stats_by_db) # digest1 in db1 should be merged self.assertEqual(merged_stats_by_db["db1"]["digest1"].count, 15) # 10 + 5 # digest2 in db2 should exist self.assertEqual(merged_stats_by_db["db2"]["digest2"].count, 3) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_asserts_n_hours(self, mock_log): """Test run asserts n_hours > 1""" begin_hour_ts = 1609459200 with self.assertRaises(AssertionError): self.task.run(begin_hour_ts, 1) @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_asserts_timestamp(self, mock_log): """Test run asserts timestamp is divisible by n_hours * 3600""" begin_hour_ts = 1609459200 n_hours = 2 invalid_ts = begin_hour_ts + 3600 with self.assertRaises(AssertionError): self.task.run(invalid_ts, n_hours) def test_upload_modified_objects(self): """Test upload_modified_objects uploads all tracked files""" # Add some files to track test_file1 = os.path.join(self.temp_dir, "file1.db") test_file2 = os.path.join(self.temp_dir, "file2.db") os.makedirs(os.path.dirname(test_file1), exist_ok=True) with open(test_file1, 'w') as f: f.write("test1") with open(test_file2, 'w') as f: f.write("test2") self.task.dirty_merged_statsdb_paths[test_file1] = "s3-key1" self.task.dirty_merged_statsdb_paths[test_file2] = "s3-key2" self.task.upload_modified_objects() # Verify upload was called for each file self.assertEqual(self.mock_s3.upload.call_count, 2) # Verify dictionary is cleared self.assertEqual(len(self.task.dirty_merged_statsdb_paths), 0) def test_upload_modified_objects_empty(self): """Test upload_modified_objects with no files""" self.task.upload_modified_objects() # Should not call upload self.mock_s3.upload.assert_not_called() @patch('o11ydiagnosis.subtasks.slowlogcompactor.task.log') def test_run_with_single_database(self, mock_log): """Test run with single database""" begin_hour_ts = 1609459200 n_hours = 2 # Still need > 1, but only one file exists utc_hour1 = datetime.utcfromtimestamp(begin_hour_ts).strftime("%Y%m%d%H") statsdb_path1 = os.path.join( self.local_path, "slowlog-stats", "1h", f"{utc_hour1}.db" ) stats1 = SlowStats("digest1") stats1.count = 10 self.create_test_statsdb( statsdb_path1, keys=["key1"], meta={"digest1": ("SELECT * FROM test", False)}, stats={"digest1": stats1} ) self.task.existed_statsdb_paths[statsdb_path1] = "s3-key1" self.task.run(begin_hour_ts, n_hours) # Should still create merged database merged_path = os.path.join( self.local_path, "slowlog-stats", f"{n_hours}h", f"{utc_hour1}.db" ) self.assertTrue(os.path.exists(merged_path)) # Verify data is preserved with sqlite3.connect(merged_path) as db: merged_keys = StatsDB.query_keys(db) merged_stats = StatsDB.query_stats(db) self.assertEqual(set(merged_keys), {"key1"}) self.assertEqual(merged_stats["digest1"].count, 10) if __name__ == '__main__': unittest.main()