| # | |
| # Licensed to the Apache Software Foundation (ASF) under one or more | |
| # contributor license agreements. See the NOTICE file distributed with | |
| # this work for additional information regarding copyright ownership. | |
| # The ASF licenses this file to You under the Apache License, Version 2.0 | |
| # (the "License"); you may not use this file except in compliance with | |
| # the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """Tests for Google Cloud Storage client.""" | |
| # pytype: skip-file | |
| import logging | |
| import os | |
| import random | |
| import unittest | |
| from datetime import datetime | |
| import mock | |
| from apache_beam import version as beam_version | |
| from apache_beam.metrics.execution import MetricsContainer | |
| from apache_beam.metrics.execution import MetricsEnvironment | |
| from apache_beam.metrics.metricbase import MetricName | |
| from apache_beam.pipeline import PipelineOptions | |
| from apache_beam.runners.worker import statesampler | |
| from apache_beam.utils import counters | |
| # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports | |
| try: | |
| from google.cloud.exceptions import BadRequest | |
| from google.cloud.exceptions import NotFound | |
| from apache_beam.io.gcp import gcsio | |
| from apache_beam.io.gcp.gcsio_retry import DEFAULT_RETRY_WITH_THROTTLING_COUNTER | |
| except ImportError: | |
| NotFound = None | |
| # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports | |
| DEFAULT_GCP_PROJECT = 'apache-beam-testing' | |
| class FakeGcsClient(object): | |
| # Fake storage client. | |
| def __init__(self): | |
| self.buckets = {} | |
| def _add_bucket(self, bucket): | |
| self.buckets[bucket.name] = bucket | |
| return self.buckets[bucket.name] | |
| def bucket(self, name): | |
| return FakeBucket(self, name) | |
| def create_bucket(self, name): | |
| return self._add_bucket(self.bucket(name)) | |
| def get_bucket(self, name): | |
| if name in self.buckets: | |
| return self.buckets[name] | |
| else: | |
| raise NotFound("Bucket not found") | |
| def lookup_bucket(self, name): | |
| if name in self.buckets: | |
| return self.buckets[name] | |
| else: | |
| return self.create_bucket(name) | |
| def batch(self, raise_exception=True): | |
| # Return a mock object configured to act as a context manager | |
| # and provide the necessary _responses attribute after __exit__. | |
| # test_delete performs 3 deletions. | |
| num_expected_responses = 3 | |
| mock_batch = mock.Mock() | |
| # Configure the mock responses (assuming success for test_delete) | |
| # These need to be available *after* the 'with' block finishes. | |
| # We'll store them temporarily and assign in __exit__. | |
| successful_responses = [ | |
| mock.Mock(status_code=204) for _ in range(num_expected_responses) | |
| ] | |
| # Define the exit logic | |
| def mock_exit_logic(exc_type, exc_val, exc_tb): | |
| # Assign responses to the mock instance itself | |
| # so they are available after the 'with' block. | |
| mock_batch._responses = successful_responses | |
| # Configure the mock to behave like a context manager | |
| mock_batch.configure_mock( | |
| __enter__=mock.Mock(return_value=mock_batch), | |
| __exit__=mock.Mock(side_effect=mock_exit_logic)) | |
| # The loop inside _batch_with_retry calls fn(request) for each item. | |
| # The real batch object might have methods like add() or similar, | |
| # but the core logic in gcsio.py calls the passed function `fn` directly | |
| # within the `with` block. So, no specific action methods seem needed | |
| # on the mock_batch itself for this test case. | |
| return mock_batch | |
| def add_file(self, bucket, blob, contents): | |
| folder = self.lookup_bucket(bucket) | |
| holder = folder.lookup_blob(blob) | |
| holder.contents = contents | |
| def get_file(self, bucket, blob): | |
| folder = self.get_bucket(bucket.name) | |
| holder = folder.get_blob(blob.name) | |
| return holder | |
| def list_blobs(self, bucket_or_path, prefix=None, **unused_kwargs): | |
| bucket = self.get_bucket(bucket_or_path.name) | |
| if not prefix: | |
| return list(bucket.blobs.values()) | |
| else: | |
| output = [] | |
| for name in list(bucket.blobs): | |
| if name[0:len(prefix)] == prefix: | |
| output.append(bucket.blobs[name]) | |
| return output | |
| class FakeBucket(object): | |
| #Fake bucket for storing test blobs locally. | |
| def __init__(self, client, name): | |
| self.client = client | |
| self.name = name | |
| self.blobs = {} | |
| self.default_kms_key_name = None | |
| def _get_canonical_bucket(self): | |
| return self.client.get_bucket(self.name) | |
| def _create_blob(self, name): | |
| return FakeBlob(name, self) | |
| def add_blob(self, blob): | |
| bucket = self._get_canonical_bucket() | |
| bucket.blobs[blob.name] = blob | |
| return bucket.blobs[blob.name] | |
| def blob(self, name): | |
| return self._create_blob(name) | |
| def copy_blob(self, blob, dest, new_name=None, **kwargs): | |
| if self.get_blob(blob.name) is None: | |
| raise NotFound("source blob not found") | |
| if not new_name: | |
| new_name = blob.name | |
| new_blob = FakeBlob(new_name, dest) | |
| dest.add_blob(new_blob) | |
| return new_blob | |
| def get_blob(self, blob_name, **unused_kwargs): | |
| bucket = self._get_canonical_bucket() | |
| if blob_name in bucket.blobs: | |
| return bucket.blobs[blob_name] | |
| else: | |
| return None | |
| def lookup_blob(self, name): | |
| bucket = self._get_canonical_bucket() | |
| if name in bucket.blobs: | |
| return bucket.blobs[name] | |
| else: | |
| new_blob = bucket._create_blob(name) | |
| return bucket.add_blob(new_blob) | |
| def set_default_kms_key_name(self, name): | |
| self.default_kms_key_name = name | |
| def delete_blob(self, name, **kwargs): | |
| bucket = self._get_canonical_bucket() | |
| if name in bucket.blobs: | |
| del bucket.blobs[name] | |
| def list_blobs(self, prefix=None, **kwargs): | |
| bucket = self._get_canonical_bucket() | |
| return self.client.list_blobs(bucket, prefix, **kwargs) | |
| class FakeBlob(object): | |
| def __init__( | |
| self, | |
| name, | |
| bucket, | |
| size=0, | |
| contents=None, | |
| generation=1, | |
| crc32c=None, | |
| kms_key_name=None, | |
| updated=None, | |
| fail_when_getting_metadata=False, | |
| fail_when_reading=False): | |
| self.name = name | |
| self.bucket = bucket | |
| self.size = size | |
| self.contents = contents | |
| self._generation = generation | |
| self.crc32c = crc32c | |
| self.kms_key_name = kms_key_name | |
| self.updated = updated | |
| self._fail_when_getting_metadata = fail_when_getting_metadata | |
| self._fail_when_reading = fail_when_reading | |
| self.generation = random.randint(0, (1 << 63) - 1) | |
| self.content_encoding = None | |
| self.content_type = None | |
| def reload(self): | |
| pass | |
| def delete(self): | |
| self.bucket.delete_blob(self.name) | |
| def download_as_bytes(self, **kwargs): | |
| blob = self.bucket.get_blob(self.name) | |
| if blob is None: | |
| raise NotFound("blob not found") | |
| return blob.contents | |
| def __eq__(self, other): | |
| return self.bucket.get_blob(self.name) is other.bucket.get_blob(other.name) | |
| def exists(self, **kwargs): | |
| return self.bucket.get_blob(self.name) is not None | |
| @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') | |
| class TestGCSPathParser(unittest.TestCase): | |
| BAD_GCS_PATHS = [ | |
| 'gs://', | |
| 'gs://bucket', | |
| 'gs:///name', | |
| 'gs:///', | |
| 'gs:/blah/bucket/name', | |
| ] | |
| def test_gcs_path(self): | |
| self.assertEqual( | |
| gcsio.parse_gcs_path('gs://bucket/name'), ('bucket', 'name')) | |
| self.assertEqual( | |
| gcsio.parse_gcs_path('gs://bucket/name/sub'), ('bucket', 'name/sub')) | |
| def test_bad_gcs_path(self): | |
| for path in self.BAD_GCS_PATHS: | |
| self.assertRaises(ValueError, gcsio.parse_gcs_path, path) | |
| self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs://bucket/') | |
| def test_gcs_path_object_optional(self): | |
| self.assertEqual( | |
| gcsio.parse_gcs_path('gs://bucket/name', object_optional=True), | |
| ('bucket', 'name')) | |
| self.assertEqual( | |
| gcsio.parse_gcs_path('gs://bucket/', object_optional=True), | |
| ('bucket', '')) | |
| def test_bad_gcs_path_object_optional(self): | |
| for path in self.BAD_GCS_PATHS: | |
| self.assertRaises(ValueError, gcsio.parse_gcs_path, path, True) | |
| class SampleOptions(object): | |
| def __init__(self, project, region, kms_key=None): | |
| self.project = DEFAULT_GCP_PROJECT | |
| self.region = region | |
| self.dataflow_kms_key = kms_key | |
| _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com" | |
| def _make_credentials(project=None, universe_domain=_DEFAULT_UNIVERSE_DOMAIN): | |
| import google.auth.credentials | |
| if project is not None: | |
| return mock.Mock( | |
| spec=google.auth.credentials.Credentials, | |
| project_id=project, | |
| universe_domain=universe_domain, | |
| ) | |
| return mock.Mock( | |
| spec=google.auth.credentials.Credentials, universe_domain=universe_domain) | |
| @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') | |
| class TestGCSIO(unittest.TestCase): | |
| def _insert_random_file( | |
| self, | |
| client, | |
| path, | |
| size=0, | |
| crc32c=None, | |
| kms_key_name=None, | |
| updated=None, | |
| fail_when_getting_metadata=False, | |
| fail_when_reading=False): | |
| bucket_name, blob_name = gcsio.parse_gcs_path(path) | |
| bucket = client.lookup_bucket(bucket_name) | |
| blob = FakeBlob( | |
| blob_name, | |
| bucket, | |
| size=size, | |
| contents=os.urandom(size), | |
| crc32c=crc32c, | |
| kms_key_name=kms_key_name, | |
| updated=updated, | |
| fail_when_getting_metadata=fail_when_getting_metadata, | |
| fail_when_reading=fail_when_reading) | |
| bucket.add_blob(blob) | |
| return blob | |
| def setUp(self): | |
| self.client = FakeGcsClient() | |
| self.gcs = gcsio.GcsIO(self.client) | |
| self.client.create_bucket("gcsio-test") | |
| def test_read_bucket_metric(self): | |
| sampler = statesampler.StateSampler('', counters.CounterFactory()) | |
| statesampler.set_current_tracker(sampler) | |
| state1 = sampler.scoped_state( | |
| 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) | |
| try: | |
| sampler.start() | |
| with state1: | |
| client = FakeGcsClient() | |
| gcs = gcsio.GcsIO(client, {"enable_bucket_read_metric_counter": True}) | |
| client.create_bucket("gcsio-test") | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| self._insert_random_file(client, file_name, file_size) | |
| reader = gcs.open(file_name, 'r') | |
| reader.read() | |
| container = MetricsEnvironment.current_container() | |
| self.assertEqual( | |
| container.get_counter( | |
| MetricName( | |
| "apache_beam.io.gcp.gcsio.BeamBlobReader", | |
| "GCS_read_bytes_counter_gcsio-test")).get_cumulative(), | |
| file_size) | |
| finally: | |
| sampler.stop() | |
| def test_write_bucket_metric(self): | |
| sampler = statesampler.StateSampler('', counters.CounterFactory()) | |
| statesampler.set_current_tracker(sampler) | |
| state1 = sampler.scoped_state( | |
| 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) | |
| try: | |
| sampler.start() | |
| with state1: | |
| client = FakeGcsClient() | |
| gcs = gcsio.GcsIO(client, {"enable_bucket_write_metric_counter": True}) | |
| client.create_bucket("gcsio-test") | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| gcsFile = gcs.open(file_name, 'w') | |
| gcsFile.write(str.encode("some text")) | |
| container = MetricsEnvironment.current_container() | |
| self.assertEqual( | |
| container.get_counter( | |
| MetricName( | |
| "apache_beam.io.gcp.gcsio.BeamBlobWriter", | |
| "GCS_write_bytes_counter_gcsio-test")).get_cumulative(), | |
| 9) | |
| finally: | |
| sampler.stop() | |
| def test_default_bucket_name(self): | |
| self.assertEqual( | |
| gcsio.default_gcs_bucket_name(DEFAULT_GCP_PROJECT, "us-central1"), | |
| 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494') | |
| def test_default_bucket_name_failure(self): | |
| self.assertEqual( | |
| gcsio.get_or_create_default_gcs_bucket( | |
| SampleOptions( | |
| DEFAULT_GCP_PROJECT, "us-central1", kms_key="kmskey!")), | |
| None) | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_ownership_match( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_bucket = mock.Mock() | |
| mock_bucket.project_number = 123456789 | |
| mock_gcsio.get_bucket.return_value = mock_bucket | |
| mock_crm_client = mock_crm_class.return_value | |
| mock_project_info = mock.Mock() | |
| mock_project_info.name = 'projects/123456789' | |
| mock_crm_client.get_project.return_value = mock_project_info | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| bucket = gcsio.get_or_create_default_gcs_bucket(options) | |
| self.assertEqual(bucket, mock_bucket) | |
| mock_gcsio.get_bucket.assert_called_once_with( | |
| 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494') | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_ownership_mismatch( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_bucket = mock.Mock() | |
| mock_bucket.project_number = 999999999 | |
| mock_gcsio.get_bucket.return_value = mock_bucket | |
| mock_crm_client = mock_crm_class.return_value | |
| mock_project_info = mock.Mock() | |
| mock_project_info.name = 'projects/123456789' | |
| mock_crm_client.get_project.return_value = mock_project_info | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| with self.assertRaises(ValueError) as ctx: | |
| gcsio.get_or_create_default_gcs_bucket(options) | |
| self.assertIn( | |
| f'is not owned by project {DEFAULT_GCP_PROJECT}', str(ctx.exception)) | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_ownership_no_bucket_number( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_bucket = mock.Mock() | |
| mock_bucket.project_number = None | |
| mock_bucket.name = 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494' | |
| mock_gcsio.get_bucket.return_value = mock_bucket | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| bucket = gcsio.get_or_create_default_gcs_bucket(options) | |
| self.assertEqual(bucket, mock_bucket) | |
| mock_crm_class.assert_not_called() | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_ownership_crm_error( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_bucket = mock.Mock() | |
| mock_bucket.project_number = 123456789 | |
| mock_gcsio.get_bucket.return_value = mock_bucket | |
| mock_crm_client = mock_crm_class.return_value | |
| mock_crm_client.get_project.side_effect = Exception("API Disabled") | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| bucket = gcsio.get_or_create_default_gcs_bucket(options) | |
| # Should fall back to success (warn only) | |
| self.assertEqual(bucket, mock_bucket) | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_ownership_mock_project_number( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| # Creating a mock bucket without setting project_number (it returns another mock object) | |
| mock_bucket = mock.Mock() | |
| mock_gcsio.get_bucket.return_value = mock_bucket | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| bucket = gcsio.get_or_create_default_gcs_bucket(options) | |
| # Verification should be skipped, returning the mock bucket and never calling CRM | |
| self.assertEqual(bucket, mock_bucket) | |
| mock_crm_class.assert_not_called() | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_conflict( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_bucket = mock.Mock() | |
| mock_bucket.project_number = 123456789 | |
| mock_gcsio.get_bucket.side_effect = [None, mock_bucket] | |
| from google.api_core.exceptions import Conflict | |
| mock_gcsio.create_bucket.side_effect = Conflict("Already owned by you") | |
| mock_crm_client = mock_crm_class.return_value | |
| mock_project_info = mock.Mock() | |
| mock_project_info.name = 'projects/123456789' | |
| mock_crm_client.get_project.return_value = mock_project_info | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| bucket = gcsio.get_or_create_default_gcs_bucket(options) | |
| self.assertEqual(bucket, mock_bucket) | |
| self.assertEqual(mock_gcsio.get_bucket.call_count, 2) | |
| mock_gcsio.create_bucket.assert_called_once() | |
| @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') | |
| @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') | |
| def test_get_or_create_default_gcs_bucket_conflict_reraise( | |
| self, mock_gcsio_class, mock_crm_class): | |
| mock_gcsio = mock_gcsio_class.return_value | |
| mock_gcsio.get_bucket.side_effect = [None, None] | |
| from google.api_core.exceptions import Conflict | |
| mock_gcsio.create_bucket.side_effect = Conflict("Bucket name unavailable") | |
| options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') | |
| with self.assertRaises(Conflict): | |
| gcsio.get_or_create_default_gcs_bucket(options) | |
| self.assertEqual(mock_gcsio.get_bucket.call_count, 2) | |
| def test_exists(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| self._insert_random_file(self.client, file_name, file_size) | |
| self.assertFalse(self.gcs.exists(file_name + 'xyz')) | |
| self.assertTrue(self.gcs.exists(file_name)) | |
| @mock.patch.object(FakeBlob, 'exists') | |
| def test_exists_failure(self, mock_get): | |
| # Raising an error other than 404. Raising 404 is a valid failure for | |
| # exists() call. | |
| mock_get.side_effect = BadRequest("Try again") | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| self._insert_random_file(self.client, file_name, file_size) | |
| with self.assertRaises(BadRequest): | |
| self.gcs.exists(file_name) | |
| def test_checksum(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| checksum = 'deadbeef' | |
| self._insert_random_file(self.client, file_name, file_size, crc32c=checksum) | |
| self.assertTrue(self.gcs.exists(file_name)) | |
| self.assertEqual(checksum, self.gcs.checksum(file_name)) | |
| def test_size(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| self._insert_random_file(self.client, file_name, file_size) | |
| self.assertTrue(self.gcs.exists(file_name)) | |
| self.assertEqual(1234, self.gcs.size(file_name)) | |
| def test_kms_key(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| kms_key_name = "dummy" | |
| self._insert_random_file( | |
| self.client, file_name, file_size, kms_key_name=kms_key_name) | |
| self.assertTrue(self.gcs.exists(file_name)) | |
| self.assertEqual(kms_key_name, self.gcs.kms_key(file_name)) | |
| def test_last_updated(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| updated = datetime.fromtimestamp(123456.78) | |
| self._insert_random_file(self.client, file_name, file_size, updated=updated) | |
| self.assertTrue(self.gcs.exists(file_name)) | |
| self.assertEqual( | |
| gcsio.GcsIO._updated_to_seconds(updated), | |
| self.gcs.last_updated(file_name)) | |
| def test_file_status(self): | |
| file_name = 'gs://gcsio-test/dummy_file' | |
| file_size = 1234 | |
| updated = datetime.fromtimestamp(123456.78) | |
| checksum = 'deadbeef' | |
| self._insert_random_file( | |
| self.client, file_name, file_size, updated=updated, crc32c=checksum) | |
| file_checksum = self.gcs.checksum(file_name) | |
| file_status = self.gcs._status(file_name) | |
| self.assertEqual(file_status['size'], file_size) | |
| self.assertEqual(file_status['checksum'], file_checksum) | |
| self.assertEqual( | |
| file_status['updated'], gcsio.GcsIO._updated_to_seconds(updated)) | |
| def test_file_mode_calls(self): | |
| file_name = 'gs://gcsio-test/dummy_mode_file' | |
| self._insert_random_file(self.client, file_name) | |
| with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobWriter') as writer: | |
| self.gcs.open(file_name, 'wb') | |
| writer.assert_called() | |
| with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobReader') as reader: | |
| self.gcs.open(file_name, 'rb') | |
| reader.assert_called() | |
| def test_bad_file_modes(self): | |
| file_name = 'gs://gcsio-test/dummy_mode_file' | |
| self._insert_random_file(self.client, file_name) | |
| with self.assertRaises(ValueError): | |
| self.gcs.open(file_name, 'w+') | |
| with self.assertRaises(ValueError): | |
| self.gcs.open(file_name, 'r+b') | |
| def test_delete(self): | |
| # File path. | |
| file_name = 'gs://gcsio-test/delete_me' | |
| file_size = 1024 | |
| bucket_name, blob_name = gcsio.parse_gcs_path(file_name) | |
| # Test deletion of non-existent file. | |
| bucket = self.client.get_bucket(bucket_name) | |
| self.gcs.delete(file_name) | |
| # Insert a random file for testing. | |
| self._insert_random_file(self.client, file_name, file_size) | |
| self.assertTrue(blob_name in bucket.blobs) | |
| # Deleting the file. | |
| self.gcs.delete(file_name) | |
| self.assertFalse(blob_name in bucket.blobs) | |
| # Now test deleting a directory (prefix) with multiple files. | |
| prefix = 'gs://gcsio-test/directory_to_delete/' | |
| file_names = [f"{prefix}file1", f"{prefix}file2", f"{prefix}file3"] | |
| blobs = [gcsio.parse_gcs_path(file_name) for file_name in file_names] | |
| # Insert random files under the prefix. | |
| for file_name in file_names: | |
| self._insert_random_file(self.client, file_name, file_size) | |
| # Verify the files exist before deletion | |
| for blob in blobs: | |
| self.assertTrue(blob[1] in bucket.blobs) | |
| # Deleting the directory (all files under the prefix). | |
| self.gcs.delete(prefix, recursive=True) | |
| # Verify that the files are deleted. | |
| for blob in blobs: | |
| self.assertFalse(blob[1] in bucket.blobs) | |
| def test_copy(self): | |
| src_file_name = 'gs://gcsio-test/source' | |
| dest_file_name = 'gs://gcsio-test/dest' | |
| src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) | |
| dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) | |
| src_bucket = self.client.lookup_bucket(src_bucket_name) | |
| dest_bucket = self.client.lookup_bucket(dest_bucket_name) | |
| file_size = 1024 | |
| self._insert_random_file(self.client, src_file_name, file_size) | |
| self.assertTrue(src_blob_name in src_bucket.blobs) | |
| self.assertFalse(dest_blob_name in dest_bucket.blobs) | |
| self.gcs.copy(src_file_name, dest_file_name) | |
| self.assertTrue(src_blob_name in src_bucket.blobs) | |
| self.assertTrue(dest_blob_name in dest_bucket.blobs) | |
| # Test copy of non-existent files. | |
| with self.assertRaises(NotFound): | |
| self.gcs.copy( | |
| 'gs://gcsio-test/non-existent', | |
| 'gs://gcsio-test/non-existent-destination') | |
| @staticmethod | |
| def _fake_batch_responses(status_codes): | |
| return mock.Mock( | |
| __enter__=mock.Mock(), | |
| __exit__=mock.Mock(), | |
| _responses=[ | |
| mock.Mock( | |
| **{ | |
| 'json.return_value': { | |
| 'error': { | |
| 'message': 'error' | |
| } | |
| }, | |
| 'request.method': 'BATCH', | |
| 'request.url': 'contentid://None', | |
| }, | |
| status_code=code, | |
| ) for code in status_codes | |
| ], | |
| ) | |
| @mock.patch('apache_beam.io.gcp.gcsio.MAX_BATCH_OPERATION_SIZE', 3) | |
| @mock.patch('time.sleep', mock.Mock()) | |
| def test_copy_batch(self): | |
| src_dest_pairs = [ | |
| (f'gs://source_bucket/file{i}.txt', f'gs://dest_bucket/file{i}.txt') | |
| for i in range(7) | |
| ] | |
| gcs_io = gcsio.GcsIO( | |
| storage_client=mock.Mock( | |
| batch=mock.Mock( | |
| side_effect=[ | |
| self._fake_batch_responses([200, 404, 429]), | |
| self._fake_batch_responses([429]), | |
| self._fake_batch_responses([429]), | |
| self._fake_batch_responses([200]), | |
| self._fake_batch_responses([200, 429, 200]), | |
| self._fake_batch_responses([200]), | |
| self._fake_batch_responses([200]), | |
| ]), | |
| )) | |
| results = gcs_io.copy_batch(src_dest_pairs) | |
| expected = [ | |
| ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt', None), | |
| ('gs://source_bucket/file1.txt', 'gs://dest_bucket/file1.txt', 404), | |
| ('gs://source_bucket/file2.txt', 'gs://dest_bucket/file2.txt', None), | |
| ('gs://source_bucket/file3.txt', 'gs://dest_bucket/file3.txt', None), | |
| ('gs://source_bucket/file4.txt', 'gs://dest_bucket/file4.txt', None), | |
| ('gs://source_bucket/file5.txt', 'gs://dest_bucket/file5.txt', None), | |
| ('gs://source_bucket/file6.txt', 'gs://dest_bucket/file6.txt', None), | |
| ] | |
| self.assertEqual(results, expected) | |
| @mock.patch('time.sleep', mock.Mock()) | |
| @mock.patch('time.monotonic', mock.Mock(side_effect=[0, 120])) | |
| def test_copy_batch_timeout_exceeded(self): | |
| src_dest_pairs = [ | |
| ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt') | |
| ] | |
| gcs_io = gcsio.GcsIO( | |
| storage_client=mock.Mock( | |
| batch=mock.Mock(side_effect=[self._fake_batch_responses([429])]))) | |
| results = gcs_io.copy_batch(src_dest_pairs) | |
| expected = [ | |
| ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt', 429), | |
| ] | |
| self.assertEqual(results, expected) | |
| def test_copytree(self): | |
| src_dir_name = 'gs://gcsio-test/source/' | |
| dest_dir_name = 'gs://gcsio-test/dest/' | |
| file_size = 1024 | |
| paths = ['a', 'b/c', 'b/d'] | |
| for path in paths: | |
| src_file_name = src_dir_name + path | |
| dest_file_name = dest_dir_name + path | |
| src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) | |
| dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) | |
| src_bucket = self.client.lookup_bucket(src_bucket_name) | |
| dest_bucket = self.client.lookup_bucket(dest_bucket_name) | |
| file_size = 1024 | |
| self._insert_random_file(self.client, src_file_name, file_size) | |
| self.assertTrue(src_blob_name in src_bucket.blobs) | |
| self.assertFalse(dest_blob_name in dest_bucket.blobs) | |
| self.gcs.copytree(src_dir_name, dest_dir_name) | |
| for path in paths: | |
| src_file_name = src_dir_name + path | |
| dest_file_name = dest_dir_name + path | |
| src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) | |
| dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) | |
| src_bucket = self.client.lookup_bucket(src_bucket_name) | |
| dest_bucket = self.client.lookup_bucket(dest_bucket_name) | |
| self.assertTrue(src_blob_name in src_bucket.blobs) | |
| self.assertTrue(dest_blob_name in dest_bucket.blobs) | |
| def test_rename(self): | |
| src_file_name = 'gs://gcsio-test/source' | |
| dest_file_name = 'gs://gcsio-test/dest' | |
| src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) | |
| dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) | |
| file_size = 1024 | |
| self._insert_random_file(self.client, src_file_name, file_size) | |
| src_bucket = self.client.lookup_bucket(src_bucket_name) | |
| dest_bucket = self.client.lookup_bucket(dest_bucket_name) | |
| self.assertTrue(src_blob_name in src_bucket.blobs) | |
| self.assertFalse(dest_blob_name in dest_bucket.blobs) | |
| self.gcs.rename(src_file_name, dest_file_name) | |
| self.assertFalse(src_blob_name in src_bucket.blobs) | |
| self.assertTrue(dest_blob_name in dest_bucket.blobs) | |
| def test_file_buffered_read_call(self): | |
| file_name = 'gs://gcsio-test/read_line_file' | |
| read_buffer_size = 1024 | |
| self._insert_random_file(self.client, file_name, 10240) | |
| bucket_name, blob_name = gcsio.parse_gcs_path(file_name) | |
| bucket = self.client.get_bucket(bucket_name) | |
| blob = bucket.get_blob(blob_name) | |
| with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobReader') as reader: | |
| self.gcs.open(file_name, read_buffer_size=read_buffer_size) | |
| reader.assert_called_with( | |
| blob, | |
| chunk_size=read_buffer_size, | |
| enable_read_bucket_metric=False, | |
| retry=DEFAULT_RETRY_WITH_THROTTLING_COUNTER) | |
| def test_file_write_call(self): | |
| file_name = 'gs://gcsio-test/write_file' | |
| with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobWriter') as writer: | |
| self.gcs.open(file_name, 'w') | |
| writer.assert_called() | |
| def test_list_files(self): | |
| bucket_name = 'gcsio-test' | |
| objects = [ | |
| ('cow/cat/fish', 2), | |
| ('cow/cat/blubber', 3), | |
| ('cow/dog/blubber', 4), | |
| ] | |
| for (object_name, size) in objects: | |
| file_name = 'gs://%s/%s' % (bucket_name, object_name) | |
| self._insert_random_file(self.client, file_name, size) | |
| test_cases = [ | |
| ( | |
| 'gs://gcsio-test/c', | |
| [ | |
| ('cow/cat/fish', 2), | |
| ('cow/cat/blubber', 3), | |
| ('cow/dog/blubber', 4), | |
| ]), | |
| ( | |
| 'gs://gcsio-test/cow/', | |
| [ | |
| ('cow/cat/fish', 2), | |
| ('cow/cat/blubber', 3), | |
| ('cow/dog/blubber', 4), | |
| ]), | |
| ('gs://gcsio-test/cow/cat/fish', [ | |
| ('cow/cat/fish', 2), | |
| ]), | |
| ] | |
| for file_pattern, expected_object_names in test_cases: | |
| expected_file_names = [('gs://%s/%s' % (bucket_name, object_name), size) | |
| for (object_name, size) in expected_object_names] | |
| self.assertEqual( | |
| set(self.gcs.list_files(file_pattern)), set(expected_file_names)) | |
| def test_downloader_fail_non_existent_object(self): | |
| file_name = 'gs://gcsio-metrics-test/dummy_mode_file' | |
| with self.assertRaises(NotFound): | |
| with self.gcs.open(file_name, 'r') as f: | |
| f.read(1) | |
| def test_blob_delete(self): | |
| file_name = 'gs://gcsio-test/delete_me' | |
| file_size = 1024 | |
| bucket_name, blob_name = gcsio.parse_gcs_path(file_name) | |
| # Test deletion of non-existent file. | |
| bucket = self.client.get_bucket(bucket_name) | |
| self.gcs.delete(file_name) | |
| self._insert_random_file(self.client, file_name, file_size) | |
| self.assertTrue(blob_name in bucket.blobs) | |
| blob = bucket.get_blob(blob_name) | |
| self.assertIsNotNone(blob) | |
| blob.delete() | |
| self.assertFalse(blob_name in bucket.blobs) | |
| @mock.patch('google.cloud._http.JSONConnection._do_request') | |
| @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials') | |
| def test_headers(self, mock_get_service_credentials, mock_do_request): | |
| from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter | |
| mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter( | |
| _make_credentials("test-project")) | |
| options = PipelineOptions([ | |
| "--job_name=test-job-name", | |
| "--gcs_custom_audit_entry=user=test-user-id", | |
| "--gcs_custom_audit_entries={\"id\": \"1234\", \"status\": \"ok\"}" | |
| ]) | |
| gcs = gcsio.GcsIO(pipeline_options=options) | |
| # no HTTP request when initializing GcsIO | |
| mock_do_request.assert_not_called() | |
| import requests | |
| response = requests.Response() | |
| response.status_code = 200 | |
| mock_do_request.return_value = response | |
| # The function of get_bucket() is supposed to send only one HTTP request | |
| gcs.get_bucket("test-bucket") | |
| mock_do_request.assert_called_once() | |
| call_args = mock_do_request.call_args[0] | |
| # Headers are specified as the third argument of | |
| # google.cloud._http.JSONConnection._do_request | |
| actual_headers = call_args[2] | |
| beam_user_agent = "apache-beam/%s (GPN:Beam)" % beam_version.__version__ | |
| self.assertIn(beam_user_agent, actual_headers['User-Agent']) | |
| self.assertEqual(actual_headers['x-goog-custom-audit-job'], 'test-job-name') | |
| self.assertEqual(actual_headers['x-goog-custom-audit-user'], 'test-user-id') | |
| self.assertEqual(actual_headers['x-goog-custom-audit-id'], '1234') | |
| self.assertEqual(actual_headers['x-goog-custom-audit-status'], 'ok') | |
| @mock.patch('google.cloud._http.JSONConnection._do_request') | |
| @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials') | |
| def test_create_default_bucket( | |
| self, mock_get_service_credentials, mock_do_request): | |
| from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter | |
| mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter( | |
| _make_credentials("test-project")) | |
| gcs = gcsio.GcsIO(pipeline_options={"job_name": "test-job-name"}) | |
| # no HTTP request when initializing GcsIO | |
| mock_do_request.assert_not_called() | |
| import requests | |
| response = requests.Response() | |
| response.status_code = 200 | |
| mock_do_request.return_value = response | |
| # The function of create_bucket() is supposed to send only one HTTP request | |
| gcs.create_bucket("test-bucket", "test-project") | |
| mock_do_request.assert_called_once() | |
| call_args = mock_do_request.call_args[0] | |
| # Request data is specified as the fourth argument of | |
| # google.cloud._http.JSONConnection._do_request | |
| actual_request_data = call_args[3] | |
| import json | |
| request_data_json = json.loads(actual_request_data) | |
| # verify soft delete policy is disabled by default in the bucket creation | |
| # request | |
| self.assertEqual( | |
| request_data_json['softDeletePolicy']['retentionDurationSeconds'], 0) | |
| @mock.patch("apache_beam.io.gcp.gcsio.GcsIO.get_bucket") | |
| def test_is_soft_delete_enabled(self, mock_get_bucket): | |
| bucket = mock.MagicMock() | |
| mock_get_bucket.return_value = bucket | |
| # soft delete policy enabled | |
| bucket.soft_delete_policy.retention_duration_seconds = 1024 | |
| self.assertTrue( | |
| self.gcs.is_soft_delete_enabled("gs://beam_with_soft_delete/tmp")) | |
| # soft delete policy disabled | |
| bucket.soft_delete_policy.retention_duration_seconds = 0 | |
| self.assertFalse( | |
| self.gcs.is_soft_delete_enabled("gs://beam_without_soft_delete/tmp")) | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |