| # | |
| # 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. | |
| # | |
| import datetime | |
| import io | |
| import json | |
| import logging | |
| import os | |
| import tempfile | |
| import unittest | |
| import fastavro | |
| import mock | |
| import apache_beam as beam | |
| from apache_beam.coders.row_coder import RowCoder | |
| from apache_beam.io.gcp.pubsub import PubsubMessage | |
| from apache_beam.testing.util import AssertThat | |
| from apache_beam.testing.util import assert_that | |
| from apache_beam.testing.util import equal_to | |
| from apache_beam.typehints import schemas as schema_utils | |
| from apache_beam.utils.timestamp import Timestamp | |
| from apache_beam.yaml.yaml_transform import YamlTransform | |
| try: | |
| import jsonschema | |
| except ImportError: | |
| jsonschema = None | |
| class FakeReadFromPubSub: | |
| def __init__( | |
| self, | |
| topic, | |
| messages, | |
| subscription=None, | |
| id_attribute=None, | |
| timestamp_attribute=None): | |
| self._topic = topic | |
| self._subscription = subscription | |
| self._messages = messages | |
| self._id_attribute = id_attribute | |
| self._timestamp_attribute = timestamp_attribute | |
| def __call__( | |
| self, | |
| *, | |
| topic, | |
| subscription, | |
| with_attributes, | |
| id_label, | |
| timestamp_attribute): | |
| assert topic == self._topic | |
| assert id_label == self._id_attribute | |
| assert timestamp_attribute == self._timestamp_attribute | |
| assert subscription == self._subscription | |
| if with_attributes: | |
| data = self._messages | |
| else: | |
| data = [x.data for x in self._messages] | |
| return beam.Create(data) | |
| class FakeWriteToPubSub: | |
| def __init__( | |
| self, topic, messages, id_attribute=None, timestamp_attribute=None): | |
| self._topic = topic | |
| self._messages = messages | |
| self._id_attribute = id_attribute | |
| self._timestamp_attribute = timestamp_attribute | |
| def __call__(self, topic, *, with_attributes, id_label, timestamp_attribute): | |
| assert topic == self._topic | |
| assert with_attributes is True | |
| assert id_label == self._id_attribute | |
| assert timestamp_attribute == self._timestamp_attribute | |
| return AssertThat(equal_to(self._messages)) | |
| @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") | |
| class YamlPubSubTest(unittest.TestCase): | |
| def test_simple_read(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), | |
| PubsubMessage(b'msg2', | |
| {'attr': 'value2'})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) | |
| def test_simple_read_string(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage('äter'.encode('utf-8'), | |
| {'attr': 'value1'}), | |
| PubsubMessage('köttbullar'.encode('utf-8'), | |
| {'attr': 'value2'})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: STRING | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([beam.Row(payload='äter'), | |
| beam.Row(payload='köttbullar')])) | |
| def test_read_with_attribute(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), | |
| PubsubMessage(b'msg2', | |
| {'attr': 'value2'})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| attributes: [attr] | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row(payload=b'msg1', attr='value1'), | |
| beam.Row(payload=b'msg2', attr='value2') | |
| ])) | |
| def test_read_with_attribute_map(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), | |
| PubsubMessage(b'msg2', | |
| {'attr': 'value2'})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| attributes_map: attrMap | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row(payload=b'msg1', attrMap={'attr': 'value1'}), | |
| beam.Row(payload=b'msg2', attrMap={'attr': 'value2'}) | |
| ])) | |
| def test_read_with_publish_time_field(self): | |
| publish_time_1 = datetime.datetime( | |
| 2018, 3, 12, 13, 37, 1, 234567, tzinfo=datetime.timezone.utc) | |
| publish_time_2 = datetime.datetime( | |
| 2018, 3, 12, 13, 38, 2, 345678, tzinfo=datetime.timezone.utc) | |
| publish_time_3 = Timestamp.from_utc_datetime( | |
| datetime.datetime( | |
| 2018, 3, 12, 13, 39, 3, 456789, tzinfo=datetime.timezone.utc)) | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}, | |
| publish_time=publish_time_1), | |
| PubsubMessage(b'msg2', {'attr': 'value2'}, | |
| publish_time=publish_time_2), | |
| PubsubMessage(b'msg3', {'attr': 'value3'}, | |
| publish_time=publish_time_3), | |
| PubsubMessage(b'msg4', | |
| {'attr': 'value4'})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| publish_time_field: publish_time | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row( | |
| payload=b'msg1', | |
| publish_time=Timestamp.from_utc_datetime(publish_time_1)), | |
| beam.Row( | |
| payload=b'msg2', | |
| publish_time=Timestamp.from_utc_datetime(publish_time_2)), | |
| beam.Row(payload=b'msg3', publish_time=publish_time_3), | |
| beam.Row(payload=b'msg4', publish_time=None) | |
| ])) | |
| def test_read_with_attributes_and_publish_time_field(self): | |
| publish_time_1 = Timestamp.from_utc_datetime( | |
| datetime.datetime( | |
| 2018, 3, 12, 13, 37, 1, 234567, tzinfo=datetime.timezone.utc)) | |
| publish_time_2 = Timestamp.from_utc_datetime( | |
| datetime.datetime( | |
| 2018, 3, 12, 13, 38, 2, 345678, tzinfo=datetime.timezone.utc)) | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}, | |
| publish_time=publish_time_1), | |
| PubsubMessage(b'msg2', {'attr': 'value2'}, | |
| publish_time=publish_time_2) | |
| ])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| attributes: [attr] | |
| attributes_map: attrMap | |
| publish_time_field: publish_time | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row( | |
| payload=b'msg1', | |
| attr='value1', | |
| attrMap={'attr': 'value1'}, | |
| publish_time=publish_time_1), | |
| beam.Row( | |
| payload=b'msg2', | |
| attr='value2', | |
| attrMap={'attr': 'value2'}, | |
| publish_time=publish_time_2) | |
| ])) | |
| def test_read_with_empty_publish_time_field(self): | |
| for publish_time_field in ('', ' '): | |
| with self.subTest(publish_time_field=publish_time_field): | |
| with self.assertRaisesRegex( | |
| ValueError, 'publish_time_field must be a non-empty field name'): | |
| with beam.Pipeline( | |
| options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| _ = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| publish_time_field: "%s" | |
| ''' % publish_time_field) | |
| def test_read_with_id_attribute(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), | |
| PubsubMessage(b'msg2', {'attr': 'value2'})], | |
| id_attribute='some_attr')): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| id_attribute: some_attr | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) | |
| _avro_schema = { | |
| 'type': 'record', | |
| 'name': 'ec', | |
| 'fields': [{ | |
| 'name': 'label', 'type': 'string' | |
| }, { | |
| 'name': 'rank', 'type': 'int' | |
| }] | |
| } | |
| def _encode_avro(self, data): | |
| buffer = io.BytesIO() | |
| fastavro.schemaless_writer(buffer, self._avro_schema, data) | |
| buffer.seek(0) | |
| return buffer.read() | |
| def test_read_avro(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch( | |
| 'apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(self._encode_avro({'label': '37a', | |
| 'rank': 1}), {}), | |
| PubsubMessage(self._encode_avro({'label': '389a', | |
| 'rank': 2}), {})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: AVRO | |
| schema: %s | |
| ''' % json.dumps(self._avro_schema)) | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row(label='37a', rank=1), # linebreak | |
| beam.Row(label='389a', rank=2) | |
| ])) | |
| def test_read_json(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage( | |
| b'{"generator": {"x": 0, "y": 0}, "rank": 1}', | |
| {'weierstrass': 'y^2+y=x^3-x', 'label': '37a'}) | |
| ])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: JSON | |
| schema: | |
| type: object | |
| properties: | |
| generator: | |
| type: object | |
| properties: | |
| x: {type: integer} | |
| y: {type: integer} | |
| rank: {type: integer} | |
| attributes: [label] | |
| attributes_map: other | |
| ''') | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row( | |
| generator=beam.Row(x=0, y=0), | |
| rank=1, | |
| label='37a', | |
| other={ | |
| 'label': '37a', 'weierstrass': 'y^2+y=x^3-x' | |
| }) | |
| ])) | |
| def test_read_json_with_error_handling(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch( | |
| 'apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub(topic='my_topic', | |
| messages=[PubsubMessage('{"some_int": 123}', | |
| attributes={}), | |
| PubsubMessage('unparsable', | |
| attributes={})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: JSON | |
| schema: | |
| type: object | |
| properties: | |
| some_int: {type: integer} | |
| error_handling: | |
| output: errors | |
| ''') | |
| assert_that( | |
| result['good'], | |
| equal_to([beam.Row(some_int=123)]), | |
| label='CheckGood') | |
| assert_that( | |
| result['errors'] | beam.Map(lambda error: error.element), | |
| equal_to(['unparsable']), | |
| label='CheckErrors') | |
| def test_read_json_without_error_handling(self): | |
| with self.assertRaises(Exception): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch( | |
| 'apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub(topic='my_topic', | |
| messages=[PubsubMessage('{"some_int": 123}', | |
| attributes={}), | |
| PubsubMessage('unparsable', | |
| attributes={})])): | |
| _ = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: JSON | |
| schema: | |
| type: object | |
| properties: | |
| some_int: {type: integer} | |
| ''') | |
| def test_read_json_with_bad_schema(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage('{"some_int": 123}', | |
| attributes={}), | |
| PubsubMessage('{"some_int": "NOT"}', | |
| attributes={})])): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: JSON | |
| schema: | |
| type: object | |
| properties: | |
| some_int: {type: integer} | |
| error_handling: | |
| output: errors | |
| ''') | |
| assert_that( | |
| result['good'], | |
| equal_to([beam.Row(some_int=123)]), | |
| label='CheckGood') | |
| assert_that( | |
| result['errors'] | beam.Map(lambda error: error.element), | |
| equal_to(['{"some_int": "NOT"}']), | |
| label='CheckErrors') | |
| def test_simple_write(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub(topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {}), | |
| PubsubMessage(b'msg2', {})])): | |
| _ = ( | |
| p | beam.Create([beam.Row(a=b'msg1'), beam.Row(a=b'msg2')]) | |
| | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| ''')) | |
| def test_write_with_attribute(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), | |
| PubsubMessage(b'msg2', | |
| {'attr': 'value2'})])): | |
| _ = ( | |
| p | beam.Create([ | |
| beam.Row(a=b'msg1', attr='value1'), | |
| beam.Row(a=b'msg2', attr='value2') | |
| ]) | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| attributes: [attr] | |
| ''')) | |
| def test_write_with_attribute_map(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub(topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', | |
| {'a': 'b'}), | |
| PubsubMessage(b'msg2', | |
| {'c': 'd'})])): | |
| _ = ( | |
| p | beam.Create([ | |
| beam.Row(a=b'msg1', attrMap={'a': 'b'}), | |
| beam.Row(a=b'msg2', attrMap={'c': 'd'}) | |
| ]) | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| attributes_map: attrMap | |
| ''')) | |
| def test_write_with_id_attribute(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub(topic='my_topic', | |
| messages=[PubsubMessage(b'msg1', {}), | |
| PubsubMessage(b'msg2', {})], | |
| id_attribute='some_attr')): | |
| _ = ( | |
| p | beam.Create([beam.Row(a=b'msg1'), beam.Row(a=b'msg2')]) | |
| | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: RAW | |
| id_attribute: some_attr | |
| ''')) | |
| def test_write_avro(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch( | |
| 'apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage(self._encode_avro({'label': '37a', | |
| 'rank': 1}), {}), | |
| PubsubMessage(self._encode_avro({'label': '389a', | |
| 'rank': 2}), {})])): | |
| _ = ( | |
| p | beam.Create( | |
| [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)]) | |
| | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: AVRO | |
| ''')) | |
| def test_write_json(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub( | |
| topic='my_topic', | |
| messages=[PubsubMessage( | |
| b'{"generator": {"x": 0, "y": 0}, "rank": 1}', | |
| {'weierstrass': 'y^2+y=x^3-x', 'label': '37a'}) | |
| ])): | |
| _ = ( | |
| p | beam.Create([ | |
| beam.Row( | |
| label='37a', | |
| generator=beam.Row(x=0, y=0), | |
| rank=1, | |
| other={'weierstrass': 'y^2+y=x^3-x'}) | |
| ]) | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: JSON | |
| attributes: [label] | |
| attributes_map: other | |
| ''')) | |
| def test_write_proto(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| data = [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)] | |
| coder = RowCoder( | |
| schema_utils.named_fields_to_schema([('label', str), ('rank', int)])) | |
| expected_messages = [PubsubMessage(coder.encode(r), {}) for r in data] | |
| with mock.patch('apache_beam.io.WriteToPubSub', | |
| FakeWriteToPubSub(topic='my_topic', | |
| messages=expected_messages)): | |
| _ = ( | |
| p | beam.Create(data) | YamlTransform( | |
| ''' | |
| type: WriteToPubSub | |
| config: | |
| topic: my_topic | |
| format: PROTO | |
| ''')) | |
| def test_read_proto(self): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| data = [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)] | |
| coder = RowCoder( | |
| schema_utils.named_fields_to_schema([('label', str), ('rank', int)])) | |
| expected_messages = [PubsubMessage(coder.encode(r), {}) for r in data] | |
| with mock.patch('apache_beam.io.ReadFromPubSub', | |
| FakeReadFromPubSub(topic='my_topic', | |
| messages=expected_messages)): | |
| result = p | YamlTransform( | |
| ''' | |
| type: ReadFromPubSub | |
| config: | |
| topic: my_topic | |
| format: PROTO | |
| schema: | |
| type: object | |
| properties: | |
| label: {type: string} | |
| rank: {type: integer} | |
| ''') | |
| assert_that(result, equal_to(data)) | |
| class YamlMatchAllTest(unittest.TestCase): | |
| def test_match_all_simple(self): | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| file1 = os.path.join(temp_dir, 'file1.txt') | |
| file2 = os.path.join(temp_dir, 'file2.txt') | |
| for f in [file1, file2]: | |
| with open(f, 'w') as fout: | |
| fout.write('data') | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create( | |
| [beam.Row(pattern=os.path.join(temp_dir, 'file*.txt'))]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| config: | |
| file_pattern: pattern | |
| ''')) | |
| paths = result | beam.Map(lambda row: row.path) | |
| assert_that(paths, equal_to([file1, file2])) | |
| def test_match_all_single_field_default(self): | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| file1 = os.path.join(temp_dir, 'file1.txt') | |
| with open(file1, 'w') as fout: | |
| fout.write('data') | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create([beam.Row(my_sole_pattern=file1)]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| ''')) | |
| paths = result | beam.Map(lambda row: row.path) | |
| assert_that(paths, equal_to([file1])) | |
| def test_match_all_multiple_fields_error(self): | |
| with self.assertRaises(Exception): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| _ = ( | |
| p | |
| | beam.Create([beam.Row(pattern='foo', other_field='bar')]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| ''')) | |
| def test_match_all_empty_match_disallow_error(self): | |
| with self.assertRaises(Exception): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| _ = ( | |
| p | |
| | beam.Create([beam.Row(pattern='does_not_exist*.txt')]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| config: | |
| empty_match_treatment: DISALLOW | |
| ''')) | |
| def test_match_all_invalid_field_error(self): | |
| with self.assertRaisesRegex( | |
| ValueError, "Field 'invalid_field' not found in input schema fields"): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| _ = ( | |
| p | |
| | beam.Create([beam.Row(pattern='foo')]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| config: | |
| file_pattern: invalid_field | |
| ''')) | |
| def test_match_all_none_timestamp(self): | |
| from apache_beam.io.filesystem import FileMetadata | |
| class MockMatchAll(beam.PTransform): | |
| def expand(self, pcoll): | |
| return pcoll.pipeline | beam.Create([ | |
| FileMetadata( | |
| path='file.txt', | |
| size_in_bytes=100, | |
| last_updated_in_seconds=None) | |
| ]) | |
| with mock.patch('apache_beam.io.fileio.MatchAll', | |
| return_value=MockMatchAll()): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create([beam.Row(pattern='file.txt')]) | |
| | YamlTransform( | |
| ''' | |
| type: MatchAll | |
| ''')) | |
| assert_that( | |
| result, | |
| equal_to([ | |
| beam.Row( | |
| path='file.txt', | |
| size_in_bytes=100, | |
| last_updated_in_seconds=None) | |
| ])) | |
| class ReadFromBigQueryTest(unittest.TestCase): | |
| def test_query_without_schema_raises(self): | |
| from apache_beam.yaml.yaml_io import read_from_bigquery | |
| with self.assertRaisesRegex(ValueError, 'schema'): | |
| read_from_bigquery(query='SELECT id FROM dataset.table') | |
| def test_table_without_schema_ok(self): | |
| import unittest.mock as mock | |
| from apache_beam.yaml.yaml_io import read_from_bigquery | |
| with mock.patch('apache_beam.yaml.yaml_io.ReadFromBigQuery') as mock_rfbq: | |
| mock_rfbq.return_value = mock.MagicMock() | |
| read_from_bigquery(table='project:dataset.table') | |
| mock_rfbq.assert_called_once() | |
| call_kwargs = mock_rfbq.call_args[1] | |
| self.assertIsNone(call_kwargs.get('query_output_schema')) | |
| def test_query_with_schema_passes_through(self): | |
| import unittest.mock as mock | |
| from apache_beam.yaml.yaml_io import read_from_bigquery | |
| schema = { | |
| 'fields': [ | |
| { | |
| 'name': 'id', 'type': 'INTEGER', 'mode': 'NULLABLE' | |
| }, | |
| ] | |
| } | |
| with mock.patch('apache_beam.yaml.yaml_io.ReadFromBigQuery') as mock_rfbq: | |
| mock_rfbq.return_value = mock.MagicMock() | |
| read_from_bigquery(query='SELECT id FROM dataset.table', schema=schema) | |
| call_kwargs = mock_rfbq.call_args[1] | |
| self.assertEqual(call_kwargs['query_output_schema'], schema) | |
| def test_query_and_table_both_raises(self): | |
| from apache_beam.yaml.yaml_io import read_from_bigquery | |
| with self.assertRaises(AssertionError): | |
| read_from_bigquery( | |
| table='project:dataset.table', | |
| query='SELECT id FROM dataset.table', | |
| schema={'fields': []}) | |
| class FakeDicomSearch(beam.PTransform): | |
| def __init__( | |
| self, buffer_size=8, max_workers=5, client=None, credential=None): | |
| self.buffer_size = buffer_size | |
| self.max_workers = max_workers | |
| def expand(self, pcoll): | |
| def do_search(element): | |
| required = [ | |
| 'project_id', | |
| 'region', | |
| 'dataset_id', | |
| 'dicom_store_id', | |
| 'search_type', | |
| ] | |
| for key in required: | |
| if key not in element: | |
| return { | |
| 'result': [], | |
| 'status': 'Must have %s in the dict.' % key, | |
| 'input': element, | |
| 'success': False, | |
| } | |
| if element['search_type'] not in ('instances', 'studies', 'series'): | |
| return { | |
| 'result': [], | |
| 'status': ( | |
| 'Search type can only be "studies", ' | |
| '"instances" or "series"'), | |
| 'input': element, | |
| 'success': False, | |
| } | |
| if element.get('project_id') == 'bad_project': | |
| return { | |
| 'result': [], | |
| 'status': 500, | |
| 'input': element, | |
| 'success': False, | |
| } | |
| params = element.get('params') or {} | |
| result = [{'PatientName': 'Alice', 'params': params}] | |
| return { | |
| 'result': result, | |
| 'status': 200, | |
| 'input': element, | |
| 'success': True, | |
| } | |
| return pcoll | beam.Map(do_search) | |
| def _patch_dicom_search(): | |
| """Install a fake dicomio module so tests do not require GCP extras.""" | |
| import sys | |
| import types | |
| dicomio_mod = types.ModuleType('apache_beam.io.gcp.healthcare.dicomio') | |
| dicomio_mod.DicomSearch = FakeDicomSearch | |
| return mock.patch.dict( | |
| sys.modules, {'apache_beam.io.gcp.healthcare.dicomio': dicomio_mod}) | |
| class YamlDicomSearchTest(unittest.TestCase): | |
| def test_dicom_search_success(self): | |
| with _patch_dicom_search(): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create([ | |
| beam.Row( | |
| project_id='proj', | |
| region='us-central1', | |
| dataset_id='dataset', | |
| dicom_store_id='store', | |
| search_type='instances', | |
| params={'PatientName': 'Alice'}) | |
| ]) | |
| | YamlTransform( | |
| ''' | |
| type: DicomSearch | |
| ''')) | |
| assert_that( | |
| result | |
| | beam.Map( | |
| lambda row: | |
| (row.status, json.loads(row.result)[0]['PatientName'])), | |
| equal_to([('200', 'Alice')])) | |
| def test_dicom_search_coerces_string_limit_offset(self): | |
| with _patch_dicom_search(): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create([ | |
| beam.Row( | |
| project_id='proj', | |
| region='us-central1', | |
| dataset_id='dataset', | |
| dicom_store_id='store', | |
| search_type='instances', | |
| params={ | |
| 'PatientName': 'Alice', 'limit': '500', 'offset': '0' | |
| }) | |
| ]) | |
| | YamlTransform( | |
| ''' | |
| type: DicomSearch | |
| ''')) | |
| assert_that( | |
| result | |
| | beam.Map( | |
| lambda row: ( | |
| row.status, json.loads(row.result)[0]['params']['limit'], | |
| json.loads(row.result)[0]['params']['offset'])), | |
| equal_to([('200', 500, 0)])) | |
| def test_dicom_search_with_error_handling(self): | |
| with _patch_dicom_search(): | |
| with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| result = ( | |
| p | |
| | beam.Create([ | |
| beam.Row( | |
| project_id='proj', | |
| region='us-central1', | |
| dataset_id='dataset', | |
| dicom_store_id='store', | |
| search_type='instances'), | |
| beam.Row( | |
| project_id='bad_project', | |
| region='us-central1', | |
| dataset_id='dataset', | |
| dicom_store_id='store', | |
| search_type='instances'), | |
| ]) | |
| | YamlTransform( | |
| ''' | |
| type: DicomSearch | |
| config: | |
| error_handling: | |
| output: errors | |
| ''')) | |
| assert_that( | |
| result['good'] | beam.Map(lambda row: row.status), | |
| equal_to(['200']), | |
| label='CheckGood') | |
| assert_that( | |
| result['errors'] | beam.Map(lambda error: error.msg), | |
| equal_to(['DicomSearch failed with status: 500']), | |
| label='CheckErrors') | |
| def test_dicom_search_without_error_handling_raises(self): | |
| with self.assertRaises(Exception): | |
| with _patch_dicom_search(): | |
| with beam.Pipeline( | |
| options=beam.options.pipeline_options.PipelineOptions( | |
| pickle_library='cloudpickle')) as p: | |
| _ = ( | |
| p | |
| | beam.Create([ | |
| beam.Row( | |
| project_id='bad_project', | |
| region='us-central1', | |
| dataset_id='dataset', | |
| dicom_store_id='store', | |
| search_type='instances') | |
| ]) | |
| | YamlTransform( | |
| ''' | |
| type: DicomSearch | |
| ''')) | |
| class YamlDeltaTest(unittest.TestCase): | |
| def test_read_from_delta(self): | |
| from apache_beam.yaml.yaml_io import read_from_delta | |
| transform = read_from_delta( | |
| table="my_table", version=5, timestamp="2026-01-01T00:00:00Z") | |
| self.assertIsInstance(transform, beam.managed.Read) | |
| self.assertEqual(transform._source, "delta") | |
| def test_read_from_delta_cdc(self): | |
| from apache_beam.yaml.yaml_io import read_from_delta_cdc | |
| transform = read_from_delta_cdc( | |
| table="my_table", | |
| start_version=0, | |
| end_version=10, | |
| include_metadata_columns=["_change_type"]) | |
| self.assertIsInstance(transform, beam.managed.Read) | |
| self.assertEqual(transform._source, "delta_cdc") | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |