|
| 1 | +# |
| 2 | +# Licensed to the Apache Software Foundation (ASF) under one or more |
| 3 | +# contributor license agreements. See the NOTICE file distributed with |
| 4 | +# this work for additional information regarding copyright ownership. |
| 5 | +# The ASF licenses this file to You under the Apache License, Version 2.0 |
| 6 | +# (the "License"); you may not use this file except in compliance with |
| 7 | +# the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, software |
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | +# See the License for the specific language governing permissions and |
| 15 | +# limitations under the License. |
| 16 | +# |
| 17 | + |
| 18 | +"""Integration tests for the cross-language MQTT IO transforms |
| 19 | +(ReadFromMqtt / WriteToMqtt), served by the messaging expansion service. |
| 20 | +
|
| 21 | +Runs against an MQTT broker (Eclipse Mosquitto) started once per test class |
| 22 | +via testcontainers. MqttIO reads are unbounded (streaming), so the end-to-end |
| 23 | +read/write test runs on the Prism portable streaming runner -- the legacy |
| 24 | +DirectRunner cannot execute an unbounded read (see the |
| 25 | +MqttReadSchemaTransformProvider description). |
| 26 | +""" |
| 27 | + |
| 28 | +import logging |
| 29 | +import threading |
| 30 | +import time |
| 31 | +import unittest |
| 32 | + |
| 33 | +import pytest |
| 34 | + |
| 35 | +import apache_beam as beam |
| 36 | +from apache_beam.options.pipeline_options import PipelineOptions |
| 37 | +from apache_beam.options.pipeline_options import StandardOptions |
| 38 | +from apache_beam.testing.test_pipeline import TestPipeline |
| 39 | +from apache_beam.typehints.row_type import RowTypeConstraint |
| 40 | + |
| 41 | +# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports |
| 42 | +try: |
| 43 | + from apache_beam.io import ReadFromMqtt |
| 44 | + from apache_beam.io import WriteToMqtt |
| 45 | +except ImportError: |
| 46 | + ReadFromMqtt = None |
| 47 | + WriteToMqtt = None |
| 48 | + |
| 49 | +try: |
| 50 | + from testcontainers.core.container import DockerContainer |
| 51 | + from testcontainers.core.waiting_utils import wait_for_logs |
| 52 | +except ImportError: |
| 53 | + DockerContainer = None |
| 54 | + |
| 55 | +NUM_RECORDS = 3 |
| 56 | +BYTES_ROW = RowTypeConstraint.from_fields([('bytes', bytes)]) |
| 57 | + |
| 58 | + |
| 59 | +@pytest.mark.uses_messaging_java_expansion_service |
| 60 | +@unittest.skipIf( |
| 61 | + DockerContainer is None, 'testcontainers package is not installed') |
| 62 | +@unittest.skipIf( |
| 63 | + ReadFromMqtt is None or WriteToMqtt is None, |
| 64 | + 'MQTT cross-language wrappers are not generated') |
| 65 | +@unittest.skipIf( |
| 66 | + TestPipeline().get_pipeline_options().view_as(StandardOptions).runner |
| 67 | + is None, |
| 68 | + 'Do not run this test on precommit suites.') |
| 69 | +@unittest.skipIf( |
| 70 | + 'Dataflow' in ( |
| 71 | + TestPipeline().get_pipeline_options().view_as(StandardOptions).runner or |
| 72 | + ''), |
| 73 | + 'The testcontainers broker is not reachable from Dataflow workers; ' |
| 74 | + 'a Dataflow variant would need a remotely hosted MQTT broker.') |
| 75 | +class CrossLanguageMqttIOTest(unittest.TestCase): |
| 76 | + @classmethod |
| 77 | + def setUpClass(cls): |
| 78 | + # The broker is expensive to spin up and tear down, so start a single |
| 79 | + # shared instance for the whole class; each test uses its own topic(s). |
| 80 | + cls.start_mqtt_container(retries=3) |
| 81 | + host = cls.broker.get_container_host_ip() |
| 82 | + port = cls.broker.get_exposed_port(1883) |
| 83 | + cls.server_uri = 'tcp://%s:%s' % (host, port) |
| 84 | + |
| 85 | + @classmethod |
| 86 | + def tearDownClass(cls): |
| 87 | + # Sometimes stopping the container raises ReadTimeout. We can ignore it |
| 88 | + # here to avoid the test failure. |
| 89 | + try: |
| 90 | + cls.broker.stop() |
| 91 | + except Exception: |
| 92 | + logging.error('Could not stop the MQTT broker container.') |
| 93 | + |
| 94 | + # Creating a container with testcontainers sometimes raises ReadTimeout |
| 95 | + # error, so retry a couple of times. |
| 96 | + @classmethod |
| 97 | + def start_mqtt_container(cls, retries): |
| 98 | + for i in range(retries): |
| 99 | + try: |
| 100 | + # /mosquitto-no-auth.conf ships with the image and enables an |
| 101 | + # anonymous listener on port 1883. |
| 102 | + cls.broker = DockerContainer('eclipse-mosquitto:2').with_command( |
| 103 | + 'mosquitto -c /mosquitto-no-auth.conf').with_exposed_ports(1883) |
| 104 | + cls.broker.start() |
| 105 | + wait_for_logs(cls.broker, 'mosquitto version .* running', timeout=30) |
| 106 | + break |
| 107 | + except Exception as e: |
| 108 | + # If start() succeeded but a later step (e.g. wait_for_logs) failed, |
| 109 | + # stop the partially started container so the next retry / the raised |
| 110 | + # error does not leak a running Docker container. |
| 111 | + try: |
| 112 | + cls.broker.stop() |
| 113 | + except Exception: |
| 114 | + pass |
| 115 | + if i == retries - 1: |
| 116 | + logging.error('Unable to initialize the MQTT broker container.') |
| 117 | + raise e |
| 118 | + |
| 119 | + def _connection_configuration(self, topic, client_id): |
| 120 | + return { |
| 121 | + 'server_uri': self.server_uri, 'topic': topic, 'client_id': client_id |
| 122 | + } |
| 123 | + |
| 124 | + def test_xlang_mqtt_write(self): |
| 125 | + topic = 'xlang-mqtt-write-topic' |
| 126 | + expected_payloads = [b'msg-%d' % i for i in range(NUM_RECORDS)] |
| 127 | + subscriber_result = {} |
| 128 | + |
| 129 | + def subscribe(): |
| 130 | + # mosquitto_sub exits after receiving NUM_RECORDS messages (-C) or |
| 131 | + # after the timeout (-W), printing one payload per line. |
| 132 | + container = self.broker.get_wrapped_container() |
| 133 | + exit_code, output = container.exec_run([ |
| 134 | + 'mosquitto_sub', |
| 135 | + '-t', |
| 136 | + topic, |
| 137 | + '-q', |
| 138 | + '1', |
| 139 | + '-C', |
| 140 | + str(NUM_RECORDS), |
| 141 | + '-W', |
| 142 | + '120' |
| 143 | + ]) |
| 144 | + subscriber_result['exit_code'] = exit_code |
| 145 | + subscriber_result['output'] = output |
| 146 | + |
| 147 | + subscriber = threading.Thread(target=subscribe, daemon=True) |
| 148 | + subscriber.start() |
| 149 | + # Give the subscriber time to connect before publishing. |
| 150 | + time.sleep(5) |
| 151 | + |
| 152 | + with TestPipeline() as p: |
| 153 | + p.not_use_test_runner_api = True |
| 154 | + _ = ( |
| 155 | + p |
| 156 | + | 'CreatePayloads' >> beam.Create(expected_payloads) |
| 157 | + | 'ToRow' >> beam.Map(lambda payload: beam.Row(bytes=payload)). |
| 158 | + with_output_types(BYTES_ROW) |
| 159 | + | 'WriteToMqtt' >> WriteToMqtt( |
| 160 | + connection_configuration=self._connection_configuration( |
| 161 | + topic, 'xlang-mqtt-write'))) |
| 162 | + |
| 163 | + subscriber.join(timeout=150) |
| 164 | + self.assertEqual(subscriber_result.get('exit_code'), 0) |
| 165 | + received = sorted(subscriber_result.get('output', b'').split()) |
| 166 | + self.assertEqual(sorted(expected_payloads), received) |
| 167 | + |
| 168 | + def test_xlang_mqtt_read_write_streaming(self): |
| 169 | + """Exercises ReadFromMqtt and WriteToMqtt end to end on the Prism portable |
| 170 | + streaming runner. MqttIO read is unbounded, which the legacy DirectRunner |
| 171 | + cannot execute, so this is the single read test: an unbounded ReadFromMqtt |
| 172 | + on a source topic feeds a WriteToMqtt on a sink topic, the result is |
| 173 | + observed with a mosquitto_sub subscriber on the sink topic, and the |
| 174 | + (never-terminating) pipeline is then cancelled. |
| 175 | +
|
| 176 | + MQTT does not retain regular messages, so the reader must already be |
| 177 | + subscribed when messages are published -- a Kafka-style sequential |
| 178 | + write-then-read would read nothing. A background publisher therefore feeds |
| 179 | + the source topic continuously while the streaming pipeline runs. |
| 180 | + """ |
| 181 | + source_topic = 'xlang-mqtt-streaming-source' |
| 182 | + sink_topic = 'xlang-mqtt-streaming-sink' |
| 183 | + stop_publishing = threading.Event() |
| 184 | + subscriber_result = {} |
| 185 | + |
| 186 | + def publish_loop(): |
| 187 | + container = self.broker.get_wrapped_container() |
| 188 | + i = 0 |
| 189 | + while not stop_publishing.is_set(): |
| 190 | + container.exec_run([ |
| 191 | + 'mosquitto_pub', '-t', source_topic, '-m', 'msg-%d' % i, '-q', '1' |
| 192 | + ]) |
| 193 | + i += 1 |
| 194 | + time.sleep(0.5) |
| 195 | + |
| 196 | + def subscribe(): |
| 197 | + container = self.broker.get_wrapped_container() |
| 198 | + exit_code, output = container.exec_run([ |
| 199 | + 'mosquitto_sub', |
| 200 | + '-t', |
| 201 | + sink_topic, |
| 202 | + '-q', |
| 203 | + '1', |
| 204 | + '-C', |
| 205 | + str(NUM_RECORDS), |
| 206 | + '-W', |
| 207 | + '180' |
| 208 | + ]) |
| 209 | + subscriber_result['exit_code'] = exit_code |
| 210 | + subscriber_result['output'] = output |
| 211 | + |
| 212 | + publisher = threading.Thread(target=publish_loop, daemon=True) |
| 213 | + subscriber = threading.Thread(target=subscribe, daemon=True) |
| 214 | + publisher.start() |
| 215 | + subscriber.start() |
| 216 | + |
| 217 | + options = PipelineOptions([ |
| 218 | + '--runner=PrismRunner', |
| 219 | + '--environment_type=LOOPBACK', |
| 220 | + '--streaming', |
| 221 | + ]) |
| 222 | + p = TestPipeline(options=options) |
| 223 | + p.not_use_test_runner_api = True |
| 224 | + _ = ( |
| 225 | + p |
| 226 | + | 'ReadFromMqtt' >> ReadFromMqtt( |
| 227 | + connection_configuration=self._connection_configuration( |
| 228 | + source_topic, 'xlang-mqtt-streaming-read')) |
| 229 | + | 'Passthrough' >> beam.Map( |
| 230 | + lambda row: beam.Row(bytes=row.bytes)).with_output_types(BYTES_ROW) |
| 231 | + | 'WriteToMqtt' >> WriteToMqtt( |
| 232 | + connection_configuration=self._connection_configuration( |
| 233 | + sink_topic, 'xlang-mqtt-streaming-write'))) |
| 234 | + result = p.run() |
| 235 | + try: |
| 236 | + # The subscriber exits once NUM_RECORDS messages flowed through the |
| 237 | + # streaming pipeline (or fails the assertions below on its timeout). |
| 238 | + subscriber.join(timeout=200) |
| 239 | + finally: |
| 240 | + stop_publishing.set() |
| 241 | + publisher.join() |
| 242 | + try: |
| 243 | + result.cancel() |
| 244 | + except Exception: # pylint: disable=broad-except |
| 245 | + # The unbounded pipeline never finishes on its own; cancellation |
| 246 | + # after the assertion data was collected is best-effort. |
| 247 | + logging.warning('Ignoring error while cancelling the pipeline.') |
| 248 | + |
| 249 | + self.assertEqual(subscriber_result.get('exit_code'), 0) |
| 250 | + payloads = subscriber_result.get('output', b'').split() |
| 251 | + self.assertEqual(NUM_RECORDS, len(payloads)) |
| 252 | + for payload in payloads: |
| 253 | + self.assertTrue( |
| 254 | + payload.startswith(b'msg-'), 'Unexpected payload: %s' % payload) |
| 255 | + |
| 256 | + |
| 257 | +if __name__ == '__main__': |
| 258 | + logging.getLogger().setLevel(logging.INFO) |
| 259 | + unittest.main() |
0 commit comments