""" Test cases for o11ydiagnosis/engine/task.py """ import os import sys import tempfile import shutil import unittest import yaml from unittest.mock import patch, MagicMock, Mock, mock_open # Add project root to path project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')) sys.path.insert(0, project_root) # Add engine directory to path so relative imports work engine_path = os.path.join(project_root, 'o11ydiagnosis', 'engine') sys.path.insert(0, engine_path) from o11ydiagnosis.engine.task import Task, Manager # Import task module for accessing running_tasks import task class TestTask(unittest.TestCase): """Test Task class""" def setUp(self): """Set up test fixtures""" # Mock boto3 and other cloud storage clients to avoid credential issues self.boto3_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') self.gcs_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.GCS') self.azure_blob_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.AzureBlob') self.mock_boto3 = self.boto3_patcher.start() self.mock_gcs = self.gcs_patcher.start() self.mock_azure_blob = self.azure_blob_patcher.start() # Mock boto3.resource for S3 initialization mock_s3_resource = MagicMock() self.mock_boto3.resource.return_value = mock_s3_resource self.temp_dir = tempfile.mkdtemp() self.cache_prefix = os.path.join(self.temp_dir, 'cache') os.makedirs(self.cache_prefix, exist_ok=True) self.mock_o11y_api = MagicMock() self.mock_remote_loader = MagicMock() self.task_info = { "meta": { "clusterID": "test-cluster", "projectID": "test-project", "tenantID": "test-tenant", "o11yClusterID": "test-o11y-cluster", "clusterType": "dedicated", "provider": "aws", "region": "us-west-2", "env": "dev" }, "interval": 3600, "subtasks": [ { "type": "slowlogconverter", "config": {} } ] } self.task = Task( self.mock_o11y_api, self.task_info, self.cache_prefix, self.mock_remote_loader ) def tearDown(self): """Clean up test fixtures""" # Stop all patchers self.boto3_patcher.stop() self.gcs_patcher.stop() self.azure_blob_patcher.stop() if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_task_initialization(self): """Test Task initialization""" self.assertEqual(self.task.cluster_id, "test-cluster") self.assertEqual(self.task.project_id, "test-project") self.assertEqual(self.task.tenant_id, "test-tenant") self.assertEqual(self.task.o11y_cluster_id, "test-o11y-cluster") self.assertFalse(self.task.serverless) # "dedicated" != "serverless" self.assertEqual(self.task.interval, 3600) self.assertEqual(self.task.provider, "aws") self.assertEqual(self.task.region, "us-west-2") self.assertEqual(self.task.env, "dev") self.assertEqual(self.task.cache_prefix, self.cache_prefix) self.assertEqual(len(self.task.subtasks), 1) def test_task_serverless(self): """Test Task with serverless cluster type""" task_info = self.task_info.copy() task_info["meta"]["clusterType"] = "serverless" task = Task( self.mock_o11y_api, task_info, self.cache_prefix, self.mock_remote_loader ) self.assertTrue(task.serverless) def test_pull_file_exists(self): """Test pull when file exists""" filelist = ['file1.txt', 'file2.txt'] self.mock_remote_loader.file_exists.return_value = True self.mock_remote_loader.get_file.return_value = 'file content' # Mock the prefix attribute self.task.prefix = self.cache_prefix result = self.task.pull(filelist) self.assertEqual(len(result), 2) self.assertTrue(all(result)) self.assertEqual(self.mock_remote_loader.file_exists.call_count, 2) self.assertEqual(self.mock_remote_loader.get_file.call_count, 2) def test_pull_file_not_exists(self): """Test pull when file does not exist""" filelist = ['file1.txt'] self.mock_remote_loader.file_exists.return_value = False self.task.prefix = self.cache_prefix result = self.task.pull(filelist) self.assertEqual(len(result), 1) self.assertFalse(result[0]) self.mock_remote_loader.get_file.assert_not_called() def test_push(self): """Test push files""" filelist = ['file1.txt', 'file2.txt'] # Create test files in cache for f in filelist: file_path = os.path.join(self.cache_prefix, f) with open(file_path, 'w') as fh: fh.write('test content') result = self.task.push(filelist) self.assertEqual(len(result), 2) self.assertTrue(all(result)) self.assertEqual(self.mock_remote_loader.upload_file.call_count, 2) def test_exist_cache(self): """Test exist_cache""" filename = 'test_file.txt' file_path = os.path.join(self.cache_prefix, filename) # File doesn't exist self.assertFalse(self.task.exist_cache(filename)) # Create file with open(file_path, 'w') as f: f.write('test') # File exists self.assertTrue(self.task.exist_cache(filename)) def test_get_cache_full_path(self): """Test get_cache_full_path""" filename = 'test_file.txt' expected_path = os.path.join(self.cache_prefix, filename) self.assertEqual(self.task.get_cache_full_path(filename), expected_path) def test_create_cache(self): """Test create_cache creates directory""" filename = 'subdir/test_file.txt' self.task.create_cache(filename) expected_dir = os.path.join(self.cache_prefix, 'subdir') self.assertTrue(os.path.exists(expected_dir)) def test_save_cache(self): """Test save_cache saves content to file""" filename = 'test_file.txt' content = 'test content' self.task.save_cache(filename, content) file_path = os.path.join(self.cache_prefix, filename) self.assertTrue(os.path.exists(file_path)) with open(file_path, 'r') as f: self.assertEqual(f.read(), content) class TestManager(unittest.TestCase): """Test Manager class""" def setUp(self): """Set up test fixtures""" # Mock boto3 and other cloud storage clients to avoid credential issues self.boto3_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') self.gcs_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.GCS') self.azure_blob_patcher = patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.AzureBlob') self.mock_boto3 = self.boto3_patcher.start() self.mock_gcs = self.gcs_patcher.start() self.mock_azure_blob = self.azure_blob_patcher.start() # Mock boto3.resource for S3 initialization mock_s3_resource = MagicMock() self.mock_boto3.resource.return_value = mock_s3_resource self.temp_dir = tempfile.mkdtemp() self.cache_root = os.path.join(self.temp_dir, 'cache') os.makedirs(self.cache_root, exist_ok=True) self.mock_engine_config = MagicMock() self.mock_engine_config.cache_root = self.cache_root self.mock_engine_config.env = 'dev' self.mock_engine_config.o11y_auth0_client_id = 'test-client-id' self.mock_engine_config.o11y_auth0_client_secret = 'test-secret' self.mock_engine_config.use_local_loader = 'True' self.mock_engine_config.start_timestemp = None self.mock_engine_config.end_timestemp = None self.mock_remote_loader = MagicMock() self.mock_remote_loader.prefix = 'test-prefix' self.mock_engine_config.remote_loader = self.mock_remote_loader self.task_file = 'test-task.yaml' self.task_info_content = """ meta: clusterID: test-cluster projectID: test-project tenantID: test-tenant o11yClusterID: test-o11y-cluster clusterType: dedicated provider: aws region: us-west-2 env: dev createTime: 1234567890 interval: 3600 subtasks: - type: slowlogconverter config: {} """ self.mock_remote_loader.get_file.return_value = self.task_info_content def tearDown(self): """Clean up test fixtures""" # Stop all patchers self.boto3_patcher.stop() self.gcs_patcher.stop() self.azure_blob_patcher.stop() if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) # Clear running_tasks task.running_tasks.clear() def test_manager_initialization(self): """Test Manager initialization""" manager = Manager(self.mock_engine_config, self.task_file) self.assertEqual(manager.running_task_key, self.task_file) self.assertIsNotNone(manager.o11y_api) self.assertIsNotNone(manager.task) self.assertEqual(manager.task.cluster_id, 'test-cluster') self.assertEqual(manager.task_info['meta']['clusterID'], 'test-cluster') def test_manager_use_local_loader(self): """Test Manager with local loader""" manager = Manager(self.mock_engine_config, f"{self.mock_remote_loader.prefix}/{self.task_file}") # Should strip prefix when use_local_loader is True self.mock_remote_loader.get_file.assert_called() call_args = self.mock_remote_loader.get_file.call_args[0][0] self.assertEqual(call_args, self.task_file) @patch('o11ydiagnosis.engine.task.time') def test_load_checkpoint_exists(self, mock_time): """Test load_checkpoint when checkpoint exists""" checkpoint_content = yaml.dump({"checkpoint_time": 1234567890}) self.mock_remote_loader.file_exists.return_value = True # get_file is called twice: once for task_info in __init__, once for checkpoint in load_checkpoint self.mock_remote_loader.get_file.side_effect = [ self.task_info_content, # First call for task_info checkpoint_content # Second call for checkpoint ] manager = Manager(self.mock_engine_config, self.task_file) checkpoint = manager.load_checkpoint() self.assertEqual(checkpoint, 1234567890) def test_load_checkpoint_not_exists(self): """Test load_checkpoint when checkpoint does not exist""" self.mock_remote_loader.file_exists.return_value = False manager = Manager(self.mock_engine_config, self.task_file) checkpoint = manager.load_checkpoint() self.assertIsNone(checkpoint) def test_save_checkpoint(self): """Test save_checkpoint""" manager = Manager(self.mock_engine_config, self.task_file) timestamp = 1234567890 manager.save_checkpoint(timestamp) # Check that checkpoint file was created in cache checkpoint_path = os.path.join(self.cache_root, 'test-cluster', 'checkpoint.yaml') self.assertTrue(os.path.exists(checkpoint_path)) # Check that upload_file was called self.mock_remote_loader.upload_file.assert_called_once() @patch('o11ydiagnosis.engine.task.time') @patch('o11ydiagnosis.engine.task.running_tasks') def test_start_with_checkpoint(self, mock_running_tasks, mock_time): """Test start method with existing checkpoint""" mock_time.time.return_value = 1234567890 mock_time.sleep = Mock() # Set up running task mock_running_tasks.__getitem__ = Mock(return_value=lambda: True) mock_running_tasks.__contains__ = Mock(return_value=True) checkpoint_content = yaml.dump({"checkpoint_time": 1234567800}) self.mock_remote_loader.file_exists.return_value = True self.mock_remote_loader.get_file.side_effect = [ self.task_info_content, # First call for task info checkpoint_content # Second call for checkpoint ] manager = Manager(self.mock_engine_config, self.task_file) # Mock subtasks to prevent actual execution manager.task.subtasks = [] manager.engine_config.end_timestemp = 1234567890 # Set end time to stop loop # This will run the loop, but we need to mock it to stop quickly with patch.object(manager, 'execute_subtasks', return_value=(True, False)): # Mock running_tasks to return False after first iteration to stop loop call_count = [0] def mock_running(): call_count[0] += 1 return call_count[0] < 2 # Return True once, then False mock_running_tasks.__getitem__ = Mock(return_value=mock_running) try: manager.start() except Exception: pass # May raise due to mocking, but we're testing the structure @patch('o11ydiagnosis.engine.task.time') def test_execute_subtasks(self, mock_time): """Test execute_subtasks""" manager = Manager(self.mock_engine_config, self.task_file) # Mock subtask mock_subtask = MagicMock() mock_subtask.is_timing.return_value = True mock_subtask.ready.return_value = True manager.task.subtasks = [mock_subtask] succeed, true_started = manager.execute_subtasks('test-cluster', 1234567890) self.assertTrue(succeed) self.assertTrue(true_started) mock_subtask.start.assert_called_once() mock_subtask.close.assert_called_once() @patch('o11ydiagnosis.engine.task.time') def test_execute_subtasks_not_timing(self, mock_time): """Test execute_subtasks when subtask is not timing""" manager = Manager(self.mock_engine_config, self.task_file) # Mock subtask that is not timing mock_subtask = MagicMock() mock_subtask.is_timing.return_value = False manager.task.subtasks = [mock_subtask] succeed, true_started = manager.execute_subtasks('test-cluster', 1234567890) self.assertTrue(succeed) self.assertFalse(true_started) mock_subtask.start.assert_not_called() mock_subtask.close.assert_not_called() @patch('o11ydiagnosis.engine.task.time') def test_execute_subtasks_exception(self, mock_time): """Test execute_subtasks handles exceptions""" manager = Manager(self.mock_engine_config, self.task_file) # Mock subtask that raises exception mock_subtask = MagicMock() mock_subtask.is_timing.return_value = True mock_subtask.ready.return_value = True mock_subtask.start.side_effect = Exception("Test error") manager.task.subtasks = [mock_subtask] succeed, true_started = manager.execute_subtasks('test-cluster', 1234567890) self.assertFalse(succeed) self.assertTrue(true_started) # Started before exception if __name__ == '__main__': unittest.main()