""" Test cases for o11ydiagnosis/subtasks/slowlogconverter/objstore.py """ import os import sys import tempfile import shutil import unittest import gzip from unittest.mock import patch, MagicMock, Mock from io import BytesIO from botocore.exceptions import ClientError # Add project root to path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../'))) from o11ydiagnosis.subtasks.slowlogconverter.objstore import ( ObjectStore, Local, S3, GCS, AzureBlob ) class TestObjectStore(unittest.TestCase): """Test base ObjectStore class""" def test_objectstore_is_abstract(self): """Test that ObjectStore is an abstract base class""" with self.assertRaises(NotImplementedError): store = ObjectStore() store.list("bucket", "prefix") with self.assertRaises(NotImplementedError): store = ObjectStore() store.object("bucket", "key") with self.assertRaises(NotImplementedError): store = ObjectStore() store.upload("bucket", "key", "filepath") with self.assertRaises(NotImplementedError): store = ObjectStore() store.download("bucket", "key", "filepath") with self.assertRaises(NotImplementedError): store = ObjectStore() store.existed("bucket", "key") class TestLocal(unittest.TestCase): """Test Local file system implementation""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.store = Local(self.temp_dir) def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_list_empty_directory(self): """Test listing files in empty directory""" keys = self.store.list("bucket", "prefix") self.assertEqual(keys, []) def test_list_files(self): """Test listing files with prefix""" # Create test files test_dir = os.path.join(self.temp_dir, "test") os.makedirs(test_dir, exist_ok=True) # Create files with open(os.path.join(test_dir, "file1.txt"), "w") as f: f.write("content1") with open(os.path.join(test_dir, "file2.txt"), "w") as f: f.write("content2") with open(os.path.join(test_dir, "other.txt"), "w") as f: f.write("content3") # List with prefix keys = self.store.list("bucket", "test/file") self.assertEqual(len(keys), 2) self.assertIn("test/file1.txt", keys) self.assertIn("test/file2.txt", keys) self.assertNotIn("test/other.txt", keys) def test_object_regular_file(self): """Test reading a regular file""" test_file = os.path.join(self.temp_dir, "test.txt") with open(test_file, "w") as f: f.write("test content") obj = self.store.object("bucket", "test.txt") self.assertIsNotNone(obj) content = obj.read() self.assertEqual(content, b"test content") obj.close() def test_object_gzip_file(self): """Test reading a gzip file""" test_file = os.path.join(self.temp_dir, "test.txt.gz") with gzip.open(test_file, "wb") as f: f.write(b"compressed content") obj = self.store.object("bucket", "test.txt.gz") self.assertIsNotNone(obj) content = obj.read() self.assertEqual(content, b"compressed content") obj.close() def test_upload_regular_file(self): """Test uploading a regular file (should be gzipped)""" source_file = os.path.join(self.temp_dir, "source.txt") with open(source_file, "w") as f: f.write("test content") self.store.upload("bucket", "dest.txt", source_file) # Check that file was uploaded and gzipped dest_file = os.path.join(self.temp_dir, "dest.txt") self.assertTrue(os.path.exists(dest_file)) with gzip.open(dest_file, "rb") as f: content = f.read() self.assertEqual(content, b"test content") def test_download_regular_file(self): """Test downloading a regular file""" source_file = os.path.join(self.temp_dir, "source.txt") with open(source_file, "w") as f: f.write("test content") dest_file = os.path.join(self.temp_dir, "dest.txt") result = self.store.download("bucket", "source.txt", dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) with open(dest_file, "r") as f: content = f.read() self.assertEqual(content, "test content") def test_existed_file_exists(self): """Test checking if an existing file exists""" test_file = os.path.join(self.temp_dir, "test.txt") with open(test_file, "w") as f: f.write("content") self.assertTrue(self.store.existed("bucket", "test.txt")) def test_list_with_empty_prefix(self): """Test listing files with empty prefix""" # Create test files with open(os.path.join(self.temp_dir, "file1.txt"), "w") as f: f.write("content1") with open(os.path.join(self.temp_dir, "file2.txt"), "w") as f: f.write("content2") keys = self.store.list("bucket", "") self.assertGreaterEqual(len(keys), 2) self.assertIn("file1.txt", keys) self.assertIn("file2.txt", keys) def test_list_with_nonexistent_directory(self): """Test listing files in nonexistent directory""" keys = self.store.list("bucket", "nonexistent/prefix") self.assertEqual(keys, []) def test_list_only_files_not_directories(self): """Test that list only returns files, not directories""" test_dir = os.path.join(self.temp_dir, "test") os.makedirs(test_dir, exist_ok=True) # Create a file and a subdirectory with open(os.path.join(test_dir, "file.txt"), "w") as f: f.write("content") os.makedirs(os.path.join(test_dir, "subdir"), exist_ok=True) keys = self.store.list("bucket", "test/") # Should only return the file, not the directory self.assertIn("test/file.txt", keys) self.assertNotIn("test/subdir", keys) def test_object_file_not_found(self): """Test reading a nonexistent file""" obj = self.store.object("bucket", "nonexistent.txt") self.assertIsNone(obj) def test_object_gzip_file_not_found(self): """Test reading a nonexistent gzip file""" obj = self.store.object("bucket", "nonexistent.txt.gz") self.assertIsNone(obj) def test_lists_multiple_prefixes(self): """Test lists method with multiple prefixes""" # Create test files test_dir1 = os.path.join(self.temp_dir, "dir1") test_dir2 = os.path.join(self.temp_dir, "dir2") os.makedirs(test_dir1, exist_ok=True) os.makedirs(test_dir2, exist_ok=True) with open(os.path.join(test_dir1, "file1.txt"), "w") as f: f.write("content1") with open(os.path.join(test_dir2, "file2.txt"), "w") as f: f.write("content2") keys = self.store.lists("bucket", ["dir1/", "dir2/"]) self.assertEqual(len(keys), 2) self.assertIn("dir1/file1.txt", keys) self.assertIn("dir2/file2.txt", keys) def test_objects_multiple_keys(self): """Test objects method with multiple keys""" # Create test files test_file1 = os.path.join(self.temp_dir, "file1.txt") test_file2 = os.path.join(self.temp_dir, "file2.txt") with open(test_file1, "w") as f: f.write("content1") with open(test_file2, "w") as f: f.write("content2") objs = self.store.objects("bucket", ["file1.txt", "file2.txt"]) self.assertEqual(len(objs), 2) self.assertIsNotNone(objs[0]) self.assertIsNotNone(objs[1]) objs[0].close() objs[1].close() def test_upload_gzip_file(self): """Test uploading a gzip file (should not be double-compressed)""" source_file = os.path.join(self.temp_dir, "source.txt.gz") original_content = b"compressed content" with gzip.open(source_file, "wb") as f: f.write(original_content) self.store.upload("bucket", "dest.txt.gz", source_file) # Check that file was uploaded without double compression # Since filepath ends with .gz, it's copied as-is (not recompressed) dest_file = os.path.join(self.temp_dir, "dest.txt.gz") self.assertTrue(os.path.exists(dest_file)) # Read the compressed content directly (it's already gzipped) with open(dest_file, "rb") as f: compressed_content = f.read() # Decompress to verify it's the same content with gzip.open(BytesIO(compressed_content)) as f: decompressed = f.read() self.assertEqual(decompressed, original_content) def test_download_gzip_file(self): """Test downloading a gzip file (should be decompressed)""" source_file = os.path.join(self.temp_dir, "source.txt.gz") with gzip.open(source_file, "wb") as f: f.write(b"compressed content") dest_file = os.path.join(self.temp_dir, "dest.txt") result = self.store.download("bucket", "source.txt.gz", dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) with open(dest_file, "rb") as f: content = f.read() self.assertEqual(content, b"compressed content") def test_download_gzip_file_without_gz_extension(self): """Test downloading a gzip file when key has .gz but filepath doesn't""" source_file = os.path.join(self.temp_dir, "source.txt.gz") with gzip.open(source_file, "wb") as f: f.write(b"compressed content") dest_file = os.path.join(self.temp_dir, "dest.txt") result = self.store.download("bucket", "source.txt.gz", dest_file) self.assertTrue(result) # Should create dest.txt (not dest.txt.gz) self.assertTrue(os.path.exists(dest_file)) self.assertFalse(os.path.exists(dest_file + ".gz")) def test_download_file_not_found(self): """Test downloading a nonexistent file""" dest_file = os.path.join(self.temp_dir, "dest.txt") result = self.store.download("bucket", "nonexistent.txt", dest_file) self.assertFalse(result) def test_existed_file_not_exists(self): """Test checking if a nonexistent file exists""" self.assertFalse(self.store.existed("bucket", "nonexistent.txt")) class TestS3(unittest.TestCase): """Test S3 implementation with mocked boto3""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) if 'AWS_ENDPOINT_S3' in os.environ: del os.environ['AWS_ENDPOINT_S3'] @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_initialization_default(self, mock_boto3): """Test S3 initialization with default config""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource s3 = S3() self.assertEqual(s3.s3, mock_s3_resource) # In slowlogconverter/objstore.py, when AWS_ENDPOINT_S3 is not set, # it calls boto3.resource("s3") without config parameter mock_boto3.resource.assert_called_once_with('s3') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_initialization_with_endpoint(self, mock_boto3): """Test S3 initialization with custom endpoint (AliCloud)""" os.environ['AWS_ENDPOINT_S3'] = 'https://oss-cn-hangzhou.aliyuncs.com' mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource s3 = S3() self.assertEqual(s3.s3, mock_s3_resource) # Should use custom endpoint and config call_args = mock_boto3.resource.call_args self.assertEqual(call_args[0][0], 's3') self.assertEqual(call_args[1]['endpoint_url'], 'https://oss-cn-hangzhou.aliyuncs.com') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_list(self, mock_boto3): """Test S3 list method""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock bucket and objects mock_bucket = MagicMock() mock_obj1 = MagicMock() mock_obj1.key = 'prefix/file1.txt' mock_obj2 = MagicMock() mock_obj2.key = 'prefix/file2.txt' mock_bucket.objects.filter.return_value = [mock_obj1, mock_obj2] mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.list('test-bucket', 'prefix/') self.assertEqual(len(result), 2) self.assertIn('prefix/file1.txt', result) self.assertIn('prefix/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_object_exists(self, mock_boto3): """Test S3 object method when object exists""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock successful get mock_bucket = MagicMock() mock_object = MagicMock() mock_body = BytesIO(b'test content') mock_object.get.return_value = {'Body': mock_body} mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.object('test-bucket', 'test-key') self.assertIsNotNone(result) self.assertEqual(result.read(), b'test content') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_object_not_found(self, mock_boto3): """Test S3 object method when object does not exist""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for 404 mock_bucket = MagicMock() mock_object = MagicMock() error_response = {'Error': {'Code': '404'}} mock_object.get.side_effect = ClientError(error_response, 'GetObject') mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.object('test-bucket', 'test-key') self.assertIsNone(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_upload(self, mock_boto3): """Test S3 upload method""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Create test file test_file = os.path.join(self.temp_dir, 'test.txt') with open(test_file, 'w') as f: f.write('test content') mock_bucket = MagicMock() mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() s3.upload('test-bucket', 'test-key', test_file) # Verify upload_fileobj was called mock_bucket.upload_fileobj.assert_called_once() @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_download(self, mock_boto3): """Test S3 download method""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock download_file mock_bucket = MagicMock() def mock_download_file(key, filepath): with open(filepath, 'wb') as f: f.write(b'test content') mock_bucket.download_file.side_effect = mock_download_file mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = s3.download('test-bucket', 'test-key', dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_existed(self, mock_boto3): """Test S3 existed method""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock successful load mock_bucket = MagicMock() mock_object = MagicMock() mock_object.load.return_value = None mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.existed('test-bucket', 'test-key') self.assertTrue(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_object_gzip(self, mock_boto3): """Test S3 object method with gzip file""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock successful get with gzip content mock_bucket = MagicMock() mock_object = MagicMock() compressed_content = gzip.compress(b'test content') mock_body = BytesIO(compressed_content) mock_object.get.return_value = {'Body': mock_body} mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.object('test-bucket', 'test-key.gz') self.assertIsNotNone(result) # Should return gzip file handle self.assertTrue(hasattr(result, 'read')) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_object_client_error_other(self, mock_boto3): """Test S3 object method with other ClientError""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for other error (not 404) mock_bucket = MagicMock() mock_object = MagicMock() error_response = {'Error': {'Code': '403'}} # Forbidden mock_object.get.side_effect = ClientError(error_response, 'GetObject') mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() with self.assertRaises(ClientError): s3.object('test-bucket', 'test-key') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_object_no_such_key(self, mock_boto3): """Test S3 object method with NoSuchKey error""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for NoSuchKey mock_bucket = MagicMock() mock_object = MagicMock() error_response = {'Error': {'Code': 'NoSuchKey'}} mock_object.get.side_effect = ClientError(error_response, 'GetObject') mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.object('test-bucket', 'test-key') self.assertIsNone(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_upload_gzip(self, mock_boto3): """Test S3 upload method with gzip key""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Create test file test_file = os.path.join(self.temp_dir, 'test.txt') with open(test_file, 'w') as f: f.write('test content') mock_bucket = MagicMock() mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() s3.upload('test-bucket', 'test-key.gz', test_file) # Verify upload_fileobj was called with ContentEncoding mock_bucket.upload_fileobj.assert_called_once() call_kwargs = mock_bucket.upload_fileobj.call_args[1] self.assertEqual(call_kwargs['ExtraArgs']['ContentEncoding'], 'gzip') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_download_gzip(self, mock_boto3): """Test S3 download method with gzip file""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock download_file with gzip content mock_bucket = MagicMock() def mock_download_file(key, filepath): with gzip.open(filepath, 'wb') as f: f.write(b'compressed content') mock_bucket.download_file.side_effect = mock_download_file mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = s3.download('test-bucket', 'test-key.gz', dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) # Should be decompressed with open(dest_file, 'rb') as f: content = f.read() self.assertEqual(content, b'compressed content') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_download_not_found(self, mock_boto3): """Test S3 download method when object does not exist""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for 404 mock_bucket = MagicMock() error_response = {'Error': {'Code': '404'}} mock_bucket.download_file.side_effect = ClientError(error_response, 'DownloadFile') mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = s3.download('test-bucket', 'test-key', dest_file) self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_download_client_error_other(self, mock_boto3): """Test S3 download method with other ClientError""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for other error (not 404) mock_bucket = MagicMock() error_response = {'Error': {'Code': '403'}} # Forbidden mock_bucket.download_file.side_effect = ClientError(error_response, 'DownloadFile') mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') with self.assertRaises(ClientError): s3.download('test-bucket', 'test-key', dest_file) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_existed_not_found(self, mock_boto3): """Test S3 existed method when object does not exist""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for 404 mock_bucket = MagicMock() mock_object = MagicMock() error_response = {'Error': {'Code': '404'}} mock_object.load.side_effect = ClientError(error_response, 'HeadObject') mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.existed('test-bucket', 'test-key') self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_existed_client_error_other(self, mock_boto3): """Test S3 existed method with other ClientError""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock ClientError for other error (not 404) mock_bucket = MagicMock() mock_object = MagicMock() error_response = {'Error': {'Code': '403'}} # Forbidden mock_object.load.side_effect = ClientError(error_response, 'HeadObject') mock_bucket.Object.return_value = mock_object mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() with self.assertRaises(ClientError): s3.existed('test-bucket', 'test-key') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_lists(self, mock_boto3): """Test S3 lists method with multiple prefixes""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock bucket and objects for different prefixes mock_bucket = MagicMock() mock_obj1 = MagicMock() mock_obj1.key = 'prefix1/file1.txt' mock_obj2 = MagicMock() mock_obj2.key = 'prefix2/file2.txt' def mock_filter(Prefix): if Prefix == 'prefix1/': return [mock_obj1] elif Prefix == 'prefix2/': return [mock_obj2] return [] mock_bucket.objects.filter.side_effect = mock_filter mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.lists('test-bucket', ['prefix1/', 'prefix2/']) self.assertEqual(len(result), 2) self.assertIn('prefix1/file1.txt', result) self.assertIn('prefix2/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.boto3') def test_s3_objects(self, mock_boto3): """Test S3 objects method with multiple keys""" mock_s3_resource = MagicMock() mock_boto3.resource.return_value = mock_s3_resource # Mock successful gets mock_bucket = MagicMock() mock_object1 = MagicMock() mock_object2 = MagicMock() mock_body1 = BytesIO(b'content1') mock_body2 = BytesIO(b'content2') mock_object1.get.return_value = {'Body': mock_body1} mock_object2.get.return_value = {'Body': mock_body2} def mock_object_getter(key): if key == 'key1': return mock_object1 elif key == 'key2': return mock_object2 return MagicMock() mock_bucket.Object.side_effect = mock_object_getter mock_s3_resource.Bucket.return_value = mock_bucket s3 = S3() result = s3.objects('test-bucket', ['key1', 'key2']) self.assertEqual(len(result), 2) self.assertIsNotNone(result[0]) self.assertIsNotNone(result[1]) class TestGCS(unittest.TestCase): """Test GCS implementation with mocked google.cloud.storage""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_initialization(self, mock_client_class): """Test GCS initialization""" mock_client = MagicMock() mock_client_class.return_value = mock_client gcs = GCS() self.assertEqual(gcs.client, mock_client) mock_client_class.assert_called_once() @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_list(self, mock_client_class): """Test GCS list method""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock blobs mock_blob1 = MagicMock() mock_blob1.name = 'prefix/file1.txt' mock_blob2 = MagicMock() mock_blob2.name = 'prefix/file2.txt' mock_client.list_blobs.return_value = [mock_blob1, mock_blob2] gcs = GCS() result = gcs.list('test-bucket', 'prefix/') self.assertEqual(len(result), 2) self.assertIn('prefix/file1.txt', result) self.assertIn('prefix/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_object_exists(self, mock_client_class): """Test GCS object method when object exists""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock bucket and blob mock_bucket = MagicMock() mock_blob = MagicMock() mock_blob.download_as_bytes.return_value = b'test content' mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() result = gcs.object('test-bucket', 'test-key') self.assertIsNotNone(result) self.assertIsInstance(result, BytesIO) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_object_not_found(self, mock_client_class): """Test GCS object method when object does not exist""" from google.cloud.exceptions import NotFound mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock NotFound exception mock_bucket = MagicMock() mock_blob = MagicMock() mock_blob.download_as_bytes.side_effect = NotFound('Object not found') mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() result = gcs.object('test-bucket', 'test-key') self.assertIsNone(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_upload(self, mock_client_class): """Test GCS upload method""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Create test file test_file = os.path.join(self.temp_dir, 'test.txt') with open(test_file, 'w') as f: f.write('test content') mock_bucket = MagicMock() mock_blob = MagicMock() mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() gcs.upload('test-bucket', 'test-key', test_file) mock_blob.upload_from_filename.assert_called_once_with(test_file) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_download(self, mock_client_class): """Test GCS download method""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock download mock_bucket = MagicMock() mock_blob = MagicMock() def mock_download_to_filename(filepath): with open(filepath, 'wb') as f: f.write(b'test content') mock_blob.download_to_filename.side_effect = mock_download_to_filename mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = gcs.download('test-bucket', 'test-key', dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_existed(self, mock_client_class): """Test GCS existed method""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock successful reload mock_bucket = MagicMock() mock_blob = MagicMock() mock_blob.reload.return_value = None mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() result = gcs.existed('test-bucket', 'test-key') self.assertTrue(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_download_not_found(self, mock_client_class): """Test GCS download method when object does not exist""" from google.cloud.exceptions import NotFound mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock NotFound exception mock_bucket = MagicMock() mock_blob = MagicMock() mock_blob.download_to_filename.side_effect = NotFound('Object not found') mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = gcs.download('test-bucket', 'test-key', dest_file) self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_existed_not_found(self, mock_client_class): """Test GCS existed method when object does not exist""" from google.cloud.exceptions import NotFound mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock NotFound exception mock_bucket = MagicMock() mock_blob = MagicMock() mock_blob.reload.side_effect = NotFound('Object not found') mock_client.get_bucket.return_value = mock_bucket with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', return_value=mock_blob): gcs = GCS() result = gcs.existed('test-bucket', 'test-key') self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_lists(self, mock_client_class): """Test GCS lists method with multiple prefixes""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock blobs for different prefixes mock_blob1 = MagicMock() mock_blob1.name = 'prefix1/file1.txt' mock_blob2 = MagicMock() mock_blob2.name = 'prefix2/file2.txt' def mock_list_blobs(bucket, prefix): if prefix == 'prefix1/': return [mock_blob1] elif prefix == 'prefix2/': return [mock_blob2] return [] mock_client.list_blobs.side_effect = mock_list_blobs gcs = GCS() result = gcs.lists('test-bucket', ['prefix1/', 'prefix2/']) self.assertEqual(len(result), 2) self.assertIn('prefix1/file1.txt', result) self.assertIn('prefix2/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Client') def test_gcs_objects(self, mock_client_class): """Test GCS objects method with multiple keys""" mock_client = MagicMock() mock_client_class.return_value = mock_client # Mock bucket and blobs mock_bucket = MagicMock() mock_client.get_bucket.return_value = mock_bucket # Create mock blobs that return BytesIO objects mock_blob1 = MagicMock() mock_blob1.download_as_bytes.return_value = b'content1' mock_blob2 = MagicMock() mock_blob2.download_as_bytes.return_value = b'content2' # Blob constructor takes (name, bucket) as arguments def mock_blob_getter(name, bucket): if name == 'key1': return mock_blob1 elif name == 'key2': return mock_blob2 return MagicMock() with patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.Blob', side_effect=mock_blob_getter): gcs = GCS() result = gcs.objects('test-bucket', ['key1', 'key2']) self.assertEqual(len(result), 2) self.assertIsNotNone(result[0]) self.assertIsNotNone(result[1]) class TestAzureBlob(unittest.TestCase): """Test AzureBlob implementation with mocked azure.storage.blob""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() os.environ['AZURE_STORAGE_ACCOUNT'] = 'test-account' def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) if 'AZURE_STORAGE_ACCOUNT' in os.environ: del os.environ['AZURE_STORAGE_ACCOUNT'] @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_initialization(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob initialization""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client azure_blob = AzureBlob() self.assertEqual(azure_blob.blob_service_client, mock_blob_service_client) mock_blob_service_client_class.assert_called_once() @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_list(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob list method""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock container client and blobs mock_container_client = MagicMock() mock_blob1 = MagicMock() mock_blob1.name = 'prefix/file1.txt' mock_blob2 = MagicMock() mock_blob2.name = 'prefix/file2.txt' mock_container_client.list_blobs.return_value = [mock_blob1, mock_blob2] mock_blob_service_client.get_container_client.return_value = mock_container_client azure_blob = AzureBlob() result = azure_blob.list('test-container', 'prefix/') self.assertEqual(len(result), 2) self.assertIn('prefix/file1.txt', result) self.assertIn('prefix/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_object_exists(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob object method when object exists""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock blob client and download mock_blob_client = MagicMock() mock_download_blob = MagicMock() stream = BytesIO() stream.write(b'test content') stream.seek(0) mock_download_blob.readinto.return_value = None mock_blob_client.download_blob.return_value = mock_download_blob mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() result = azure_blob.object('test-container', 'test-key') self.assertIsNotNone(result) self.assertIsInstance(result, BytesIO) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_upload(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob upload method""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Create test file test_file = os.path.join(self.temp_dir, 'test.txt') with open(test_file, 'w') as f: f.write('test content') mock_blob_client = MagicMock() mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() azure_blob.upload('test-container', 'test-key', test_file) mock_blob_client.upload_blob.assert_called_once() @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_download(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob download method""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock download mock_blob_client = MagicMock() mock_download_blob = MagicMock() def mock_readinto(file): file.write(b'test content') mock_download_blob.readinto.side_effect = mock_readinto mock_blob_client.download_blob.return_value = mock_download_blob mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = azure_blob.download('test-container', 'test-key', dest_file) self.assertTrue(result) self.assertTrue(os.path.exists(dest_file)) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_existed(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob existed method""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock successful get_blob_properties mock_blob_client = MagicMock() mock_blob_client.get_blob_properties.return_value = MagicMock() mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() result = azure_blob.existed('test-container', 'test-key') self.assertTrue(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_object_not_found(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob object method when object does not exist""" from azure.core.exceptions import ResourceNotFoundError mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock ResourceNotFoundError mock_blob_client = MagicMock() mock_blob_client.download_blob.side_effect = ResourceNotFoundError('Blob not found') mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() result = azure_blob.object('test-container', 'test-key') self.assertIsNone(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_download_not_found(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob download method when object does not exist""" from azure.core.exceptions import ResourceNotFoundError mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock ResourceNotFoundError mock_blob_client = MagicMock() mock_blob_client.download_blob.side_effect = ResourceNotFoundError('Blob not found') mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() dest_file = os.path.join(self.temp_dir, 'downloaded.txt') result = azure_blob.download('test-container', 'test-key', dest_file) self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_existed_not_found(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob existed method when object does not exist""" from azure.core.exceptions import ResourceNotFoundError mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock ResourceNotFoundError mock_blob_client = MagicMock() mock_blob_client.get_blob_properties.side_effect = ResourceNotFoundError('Blob not found') mock_blob_service_client.get_blob_client.return_value = mock_blob_client azure_blob = AzureBlob() result = azure_blob.existed('test-container', 'test-key') self.assertFalse(result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_lists(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob lists method with multiple prefixes""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock container client and blobs for different prefixes mock_container_client = MagicMock() mock_blob1 = MagicMock() mock_blob1.name = 'prefix1/file1.txt' mock_blob2 = MagicMock() mock_blob2.name = 'prefix2/file2.txt' def mock_list_blobs(name_starts_with): if name_starts_with == 'prefix1/': return [mock_blob1] elif name_starts_with == 'prefix2/': return [mock_blob2] return [] mock_container_client.list_blobs.side_effect = mock_list_blobs mock_blob_service_client.get_container_client.return_value = mock_container_client azure_blob = AzureBlob() result = azure_blob.lists('test-container', ['prefix1/', 'prefix2/']) self.assertEqual(len(result), 2) self.assertIn('prefix1/file1.txt', result) self.assertIn('prefix2/file2.txt', result) @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.DefaultAzureCredential') @patch('o11ydiagnosis.subtasks.slowlogconverter.objstore.BlobServiceClient') def test_azure_blob_objects(self, mock_blob_service_client_class, mock_credential_class): """Test AzureBlob objects method with multiple keys""" mock_credential = MagicMock() mock_credential_class.return_value = mock_credential mock_blob_service_client = MagicMock() mock_blob_service_client_class.return_value = mock_blob_service_client # Mock blob clients and downloads mock_blob_client1 = MagicMock() mock_blob_client2 = MagicMock() stream1 = BytesIO() stream1.write(b'content1') stream1.seek(0) stream2 = BytesIO() stream2.write(b'content2') stream2.seek(0) mock_download_blob1 = MagicMock() mock_download_blob1.readinto.return_value = None mock_download_blob1.readinto.side_effect = lambda s: s.write(b'content1') or None mock_download_blob2 = MagicMock() mock_download_blob2.readinto.return_value = None mock_download_blob2.readinto.side_effect = lambda s: s.write(b'content2') or None mock_blob_client1.download_blob.return_value = mock_download_blob1 mock_blob_client2.download_blob.return_value = mock_download_blob2 def mock_get_blob_client(container, blob): if blob == 'key1': return mock_blob_client1 elif blob == 'key2': return mock_blob_client2 return MagicMock() mock_blob_service_client.get_blob_client.side_effect = mock_get_blob_client azure_blob = AzureBlob() result = azure_blob.objects('test-container', ['key1', 'key2']) self.assertEqual(len(result), 2) self.assertIsNotNone(result[0]) self.assertIsNotNone(result[1]) class TestObjectStoreListsObjects(unittest.TestCase): """Test ObjectStore base class lists and objects methods""" def setUp(self): """Set up test fixtures""" self.temp_dir = tempfile.mkdtemp() self.store = Local(self.temp_dir) def tearDown(self): """Clean up test fixtures""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) def test_lists_max_threads(self): """Test lists method with more than S3_MAX_FETCH_THREADS prefixes""" from o11ydiagnosis.subtasks.slowlogconverter.objstore import S3_MAX_FETCH_THREADS # Create many prefixes prefixes = [f"prefix{i}/" for i in range(S3_MAX_FETCH_THREADS + 10)] # Should not raise exception result = self.store.lists("bucket", prefixes) self.assertIsInstance(result, list) def test_objects_max_threads(self): """Test objects method with more than S3_MAX_FETCH_THREADS keys""" from o11ydiagnosis.subtasks.slowlogconverter.objstore import S3_MAX_FETCH_THREADS # Create many keys keys = [f"key{i}.txt" for i in range(S3_MAX_FETCH_THREADS + 10)] # Should not raise exception result = self.store.objects("bucket", keys) self.assertIsInstance(result, list) self.assertEqual(len(result), len(keys)) if __name__ == '__main__': unittest.main()