""" Test cases for o11ydiagnosis/engine/engine.py """ import os import sys import tempfile import shutil import unittest from unittest.mock import patch, MagicMock, Mock import time # 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.engine import ( EngineConfig, compare, get_task_list, func_alive, func_not_alive, ) # Import task module after adding path import task as engine_task class TestEngineConfig(unittest.TestCase): """Test EngineConfig class""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() # Set required environment variables os.environ['ENGINE_ID'] = 'test-engine' os.environ['CACHE_ROOT'] = self.temp_dir os.environ['PROVIDER'] = 'aws' os.environ['USE_LOCAL_LOADER'] = 'True' os.environ['LOCAL_LOADER_PATH'] = self.temp_dir def tearDown(self): """Clean up test fixtures""" # Clean up environment variables env_vars_to_remove = [ 'ENGINE_ID', 'BUCKET_NAME', 'CACHE_ROOT', 'O11Y_AUTH0_CLIENT_ID', 'O11Y_AUTH0_CLIENT_SECRET', 'ENV', 'START_TIMESTEMP', 'END_TIMESTEMP', 'PROVIDER', 'GCP_PROJECT', 'USE_LOCAL_LOADER', 'LOCAL_LOADER_PATH', 'REMOTE_PROVIDER', 'AZURE_STORAGE_ACCOUNT', 'POD_NUMBER', 'TOTAL_REPLICAS' ] for var in env_vars_to_remove: if var in os.environ: del os.environ[var] if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_engine_config_defaults(self): """Test EngineConfig with default values""" config = EngineConfig() self.assertEqual(config.engine_id, 'test-engine') self.assertEqual(config.cache_root, self.temp_dir) self.assertEqual(config.env, 'dev') self.assertEqual(config.provider, 'aws') self.assertEqual(config.use_local_loader, 'True') self.assertIsNotNone(config.remote_loader) self.assertEqual(config.pod_number, 0) self.assertEqual(config.total_replicas, 1) def test_engine_config_with_all_env_vars(self): """Test EngineConfig with all environment variables set""" os.environ['BUCKET_NAME'] = 'test-bucket' os.environ['O11Y_AUTH0_CLIENT_ID'] = 'test-client-id' os.environ['O11Y_AUTH0_CLIENT_SECRET'] = 'test-client-secret' os.environ['ENV'] = 'prod' os.environ['START_TIMESTEMP'] = '1234567890' os.environ['END_TIMESTEMP'] = '1234567999' os.environ['GCP_PROJECT'] = 'test-project' os.environ['POD_NUMBER'] = 'o11ydiagnosis-deltalake-2' os.environ['TOTAL_REPLICAS'] = '3' config = EngineConfig() self.assertEqual(config.bucket_name, 'test-bucket') self.assertEqual(config.o11y_auth0_client_id, 'test-client-id') self.assertEqual(config.o11y_auth0_client_secret, 'test-client-secret') self.assertEqual(config.env, 'prod') self.assertEqual(config.start_timestemp, 1234567890) self.assertEqual(config.end_timestemp, 1234567999) self.assertEqual(config.gcp_project, 'test-project') self.assertEqual(config.pod_number, 2) # Extracted from pod name self.assertEqual(config.total_replicas, 3) def test_get_remote_loader_local(self): """Test get_remote_loader returns LocalFileLoader when USE_LOCAL_LOADER is True""" config = EngineConfig() loader = config.get_remote_loader() # Should be LocalFileLoader when USE_LOCAL_LOADER is True from diagnosis.dataloader.localloader import LocalFileLoader self.assertIsInstance(loader, LocalFileLoader) @patch.dict(os.environ, {'USE_LOCAL_LOADER': 'False', 'REMOTE_PROVIDER': 'aws', 'BUCKET_NAME': 'test-bucket'}) @patch('diagnosis.dataloader.s3loader.boto3') def test_get_remote_loader_aws(self, mock_boto3): """Test get_remote_loader returns S3Loader for AWS""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource config = EngineConfig() loader = config.get_remote_loader() from diagnosis.dataloader.s3loader import S3Loader self.assertIsInstance(loader, S3Loader) @patch.dict(os.environ, {'USE_LOCAL_LOADER': 'False', 'REMOTE_PROVIDER': 'gcp', 'BUCKET_NAME': 'test-bucket', 'GCP_PROJECT': 'test-project'}) @patch('google.cloud.storage.Client') def test_get_remote_loader_gcp(self, mock_client): """Test get_remote_loader returns GCSLoader for GCP""" config = EngineConfig() loader = config.get_remote_loader() from diagnosis.dataloader.gcsloader import GCSLoader self.assertIsInstance(loader, GCSLoader) @patch.dict(os.environ, {'USE_LOCAL_LOADER': 'False', 'REMOTE_PROVIDER': 'azure', 'BUCKET_NAME': 'test-bucket', 'AZURE_STORAGE_ACCOUNT': 'test-account'}) @patch('azure.identity.DefaultAzureCredential') def test_get_remote_loader_azure(self, mock_credential): """Test get_remote_loader returns AzureBlobLoader for Azure""" config = EngineConfig() loader = config.get_remote_loader() from diagnosis.dataloader.azureblobloader import AzureBlobLoader self.assertIsInstance(loader, AzureBlobLoader) def test_pod_number_extraction(self): """Test pod number extraction from pod name""" os.environ['POD_NUMBER'] = 'o11ydiagnosis-deltalake-5' config = EngineConfig() self.assertEqual(config.pod_number, 5) os.environ['POD_NUMBER'] = 'test-pod-10' config = EngineConfig() self.assertEqual(config.pod_number, 10) # Test with empty pod name if 'POD_NUMBER' in os.environ: del os.environ['POD_NUMBER'] config = EngineConfig() self.assertEqual(config.pod_number, 0) class TestCompareFunction(unittest.TestCase): """Test compare function""" def setUp(self): """Set up test fixtures""" # Clear running_tasks engine_task.running_tasks.clear() def tearDown(self): """Clean up test fixtures""" engine_task.running_tasks.clear() def test_compare_add_new_tasks(self): """Test compare adds new tasks""" task_list = ['task1.yaml', 'task2.yaml'] to_add, to_remove = compare(task_list) self.assertEqual(len(to_add), 2) self.assertIn('task1.yaml', to_add) self.assertIn('task2.yaml', to_add) self.assertEqual(len(to_remove), 0) def test_compare_remove_old_tasks(self): """Test compare removes old tasks""" engine_task.running_tasks['task1.yaml'] = func_alive engine_task.running_tasks['task2.yaml'] = func_alive task_list = [] to_add, to_remove = compare(task_list) self.assertEqual(len(to_add), 0) self.assertEqual(len(to_remove), 2) self.assertIn('task1.yaml', to_remove) self.assertIn('task2.yaml', to_remove) def test_compare_mixed(self): """Test compare with mixed add and remove""" engine_task.running_tasks['task1.yaml'] = func_alive engine_task.running_tasks['task2.yaml'] = func_alive task_list = ['task2.yaml', 'task3.yaml'] to_add, to_remove = compare(task_list) self.assertEqual(len(to_add), 1) self.assertIn('task3.yaml', to_add) self.assertEqual(len(to_remove), 1) self.assertIn('task1.yaml', to_remove) def test_compare_no_changes(self): """Test compare with no changes""" engine_task.running_tasks['task1.yaml'] = func_alive engine_task.running_tasks['task2.yaml'] = func_alive task_list = ['task1.yaml', 'task2.yaml'] to_add, to_remove = compare(task_list) self.assertEqual(len(to_add), 0) self.assertEqual(len(to_remove), 0) class TestGetTaskList(unittest.TestCase): """Test get_task_list function""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() os.environ['ENGINE_ID'] = 'test-engine' os.environ['CACHE_ROOT'] = self.temp_dir os.environ['PROVIDER'] = 'aws' os.environ['USE_LOCAL_LOADER'] = 'True' os.environ['LOCAL_LOADER_PATH'] = self.temp_dir os.environ['TOTAL_REPLICAS'] = '1' os.environ['POD_NUMBER'] = '' def tearDown(self): """Clean up test fixtures""" env_vars_to_remove = ['ENGINE_ID', 'CACHE_ROOT', 'PROVIDER', 'USE_LOCAL_LOADER', 'LOCAL_LOADER_PATH', 'TOTAL_REPLICAS', 'POD_NUMBER'] for var in env_vars_to_remove: if var in os.environ: del os.environ[var] if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_get_task_list_success(self): """Test get_task_list successfully retrieves task list""" config = EngineConfig() # Create test task files task_dir = os.path.join(self.temp_dir, 'diagnosis', 'tasks', 'processor-test-engine') os.makedirs(task_dir, exist_ok=True) with open(os.path.join(task_dir, 'task-001.yaml'), 'w') as f: f.write('test content') with open(os.path.join(task_dir, 'task-002.yaml'), 'w') as f: f.write('test content') with open(os.path.join(task_dir, 'not-a-task.txt'), 'w') as f: f.write('not a task') task_list = get_task_list(config) # Should only return .yaml files self.assertEqual(len(task_list), 2) self.assertTrue(all(t.endswith('.yaml') for t in task_list)) def test_get_task_list_filters_by_pod_number(self): """Test get_task_list filters tasks by pod number""" config = EngineConfig() config.total_replicas = 3 config.pod_number = 1 # Mock get_file_list to return task file names # The function extracts task number from t[-8:-5], so we need filenames like "task-001.yaml" task_files = ['task-000.yaml', 'task-001.yaml', 'task-002.yaml', 'task-003.yaml', 'task-004.yaml'] config.remote_loader.get_file_list = Mock(return_value=task_files) task_list = get_task_list(config) # Pod 1 should get tasks where task_num % 3 == 1 # Task numbers: 0, 1, 2, 3, 4 # 0 % 3 = 0, 1 % 3 = 1, 2 % 3 = 2, 3 % 3 = 0, 4 % 3 = 1 # So pod 1 should get tasks 1 and 4 expected_tasks = ['task-001.yaml', 'task-004.yaml'] self.assertEqual(len(task_list), 2) for task in task_list: self.assertIn(task, expected_tasks) def test_get_task_list_exception_handling(self): """Test get_task_list handles exceptions gracefully""" config = EngineConfig() # Mock remote_loader to raise exception with patch.object(config.remote_loader, 'get_file_list', side_effect=Exception("Test error")): task_list = get_task_list(config) self.assertEqual(task_list, []) class TestHelperFunctions(unittest.TestCase): """Test helper functions""" def test_func_alive(self): """Test func_alive returns True""" self.assertTrue(func_alive()) def test_func_not_alive(self): """Test func_not_alive returns False""" self.assertFalse(func_not_alive()) if __name__ == '__main__': unittest.main()