| #!/usr/bin/env python |
| # -*- encoding: utf-8 -*- |
| |
| # 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. |
| |
| |
| # pylint: disable=missing-docstring |
| # pylint: disable=protected-access |
| import os |
| import unittest |
| |
| from heronpy.api.serializer import default_serializer |
| from heronpy.api.topology import Topology, TopologyBuilder, TopologyType |
| from heronpy.api.stream import Stream, Grouping |
| from heronpy.api.component.component_spec import HeronComponentSpec |
| from heronpy.proto import topology_pb2 |
| |
| # required environment variable |
| # note that this test doesn't write anything to /tmp directory |
| heron_options = "cmdline.topologydefn.tmpdirectory=/tmp,cmdline.topology.initial.state=RUNNING" |
| os.environ["HERON_OPTIONS"] = heron_options |
| |
| class TestSane(Topology): |
| config = {"topology.wide.config.1": "value", |
| "spout.overriden.config": True} |
| spout = HeronComponentSpec(None, "sp_class", True, 3, inputs=None, |
| outputs=["word", "count", |
| Stream(fields=['error_msg'], name='error_stream')], |
| config={"spout.specific.config.1": "value", |
| "spout.specific.config.2": True, |
| "spout.specific.config.3": -12.4, |
| "spout.specific.config.4": [1, 2, 3], |
| "spout.overriden.config": False}) |
| bolt = HeronComponentSpec(None, "bl_class", False, 4, |
| inputs={spout: Grouping.SHUFFLE, spout['error_stream']: Grouping.ALL}) |
| |
| # pylint: disable=no-member |
| # pylint: disable=too-many-branches |
| # pylint: disable=too-many-statements |
| class TopologyTest(unittest.TestCase): |
| def setUp(self): |
| os.environ["HERON_OPTIONS"] = heron_options |
| |
| def tearDown(self): |
| os.environ.pop("HERON_OPTIONS", None) |
| |
| def test_sane_topology(self): |
| self.assertEqual(TestSane.topology_name, "TestSane") |
| |
| # topology-wide config |
| expecting_topo_config = TopologyType.DEFAULT_TOPOLOGY_CONFIG |
| expecting_topo_config.update({"topology.wide.config.1": "value", |
| "spout.overriden.config": "true"}) |
| self.assertEqual(TestSane._topo_config, expecting_topo_config) |
| |
| self.assertEqual(len(TestSane._protobuf_bolts), 1) |
| self.assertEqual(len(TestSane._protobuf_spouts), 1) |
| |
| self.assertEqual(len(TestSane._heron_specs), 2) |
| for spec in TestSane._heron_specs: |
| if spec.is_spout: |
| self.assertEqual(spec.name, "spout") |
| self.assertEqual(spec.python_class_path, "sp_class") |
| self.assertEqual(spec.parallelism, 3) |
| else: |
| self.assertEqual(spec.name, "bolt") |
| self.assertEqual(spec.python_class_path, "bl_class") |
| self.assertEqual(spec.parallelism, 4) |
| |
| self.assertTrue(isinstance(TestSane.protobuf_topology, topology_pb2.Topology)) |
| |
| proto_topo = TestSane.protobuf_topology |
| |
| ### spout protobuf ### |
| self.assertEqual(len(proto_topo.spouts), 1) |
| spout = proto_topo.spouts[0] |
| self.assertEqual(spout.comp.name, "spout") |
| self.assertEqual(spout.comp.spec, topology_pb2.ComponentObjectSpec.Value("PYTHON_CLASS_NAME")) |
| self.assertEqual(spout.comp.class_name, "sp_class") |
| expecting_spout_config = {"topology.component.parallelism": "3", |
| "spout.specific.config.1": "value", |
| "spout.specific.config.2": "true", |
| "spout.specific.config.3": "-12.4", |
| "spout.specific.config.4": default_serializer.serialize([1, 2, 3]), |
| "spout.overriden.config": "false"} |
| self.assertEqual(len(spout.comp.config.kvs), len(expecting_spout_config)) |
| for conf in spout.comp.config.kvs: |
| value = expecting_spout_config[conf.key] |
| if conf.type == topology_pb2.ConfigValueType.Value("STRING_VALUE"): |
| self.assertEqual(value, conf.value) |
| elif conf.type == topology_pb2.ConfigValueType.Value("PYTHON_SERIALIZED_VALUE"): |
| self.assertEqual(value, conf.serialized_value) |
| else: |
| self.fail() |
| |
| # output stream |
| self.assertEqual(len(spout.outputs), 2) |
| for out_stream in spout.outputs: |
| if out_stream.stream.id == "default": |
| self.assertEqual(out_stream.stream.component_name, "spout") |
| self.assertEqual(len(out_stream.schema.keys), 2) |
| else: |
| self.assertEqual(out_stream.stream.id, "error_stream") |
| self.assertEqual(out_stream.stream.component_name, "spout") |
| self.assertEqual(len(out_stream.schema.keys), 1) |
| |
| |
| ### bolt protobuf ### |
| self.assertEqual(len(proto_topo.bolts), 1) |
| bolt = proto_topo.bolts[0] |
| self.assertEqual(bolt.comp.name, "bolt") |
| self.assertEqual(bolt.comp.spec, topology_pb2.ComponentObjectSpec.Value("PYTHON_CLASS_NAME")) |
| self.assertEqual(bolt.comp.class_name, "bl_class") |
| expecting_bolt_config = {"topology.component.parallelism": "4"} |
| self.assertEqual(len(bolt.comp.config.kvs), len(expecting_bolt_config)) |
| conf = bolt.comp.config.kvs[0] |
| self.assertEqual(conf.type, topology_pb2.ConfigValueType.Value("STRING_VALUE")) |
| self.assertEqual(conf.value, expecting_bolt_config[conf.key]) |
| |
| # out stream |
| self.assertEqual(len(bolt.outputs), 0) |
| |
| # in stream |
| self.assertEqual(len(bolt.inputs), 2) |
| for in_stream in bolt.inputs: |
| if in_stream.stream.id == "default": |
| self.assertEqual(in_stream.stream.component_name, "spout") |
| self.assertEqual(in_stream.gtype, topology_pb2.Grouping.Value("SHUFFLE")) |
| else: |
| self.assertEqual(in_stream.stream.id, "error_stream") |
| self.assertEqual(in_stream.stream.component_name, "spout") |
| self.assertEqual(in_stream.gtype, topology_pb2.Grouping.Value("ALL")) |
| |
| self.assertEqual(proto_topo.state, topology_pb2.TopologyState.Value("RUNNING")) |
| |
| def test_no_spout(self): |
| with self.assertRaises(ValueError): |
| # pylint:disable = unused-variable |
| class JustBolt(Topology): |
| bolt = HeronComponentSpec(None, "bl_class", False, 4) |
| |
| def test_class_dict_to_specs(self): |
| # duplicate component name |
| class_dict = {"spout": HeronComponentSpec("same_name", "sp_cls", True, 1), |
| "bolt": HeronComponentSpec("same_name", "bl_cls", False, 2)} |
| with self.assertRaises(ValueError): |
| TopologyType.class_dict_to_specs(class_dict) |
| |
| def test_add_spout_specs(self): |
| # spout with no output |
| spec = HeronComponentSpec("spout", "sp_cls", True, 1) |
| with self.assertRaises(ValueError): |
| TopologyType.add_spout_specs(spec, {}) |
| |
| def test_add_bolt_specs(self): |
| spec = HeronComponentSpec("bolt", "bl_cls", False, 1) |
| with self.assertRaises(ValueError): |
| TopologyType.add_bolt_specs(spec, {}) |
| |
| def test_sanitize_config(self): |
| # non-string key |
| with self.assertRaises(TypeError): |
| TopologyType._sanitize_config({['k', 'e', 'y']: "value"}) |
| with self.assertRaises(TypeError): |
| TopologyType._sanitize_config({None: "value"}) |
| |
| # convert boolean value |
| ret = TopologyType._sanitize_config({"key": True}) |
| self.assertEqual(ret["key"], "true") |
| ret = TopologyType._sanitize_config({"key": False}) |
| self.assertEqual(ret["key"], "false") |
| |
| # convert int and float |
| ret = TopologyType._sanitize_config({"key": 10}) |
| self.assertEqual(ret["key"], "10") |
| ret = TopologyType._sanitize_config({"key": -2400000}) |
| self.assertEqual(ret["key"], "-2400000") |
| ret = TopologyType._sanitize_config({"key": 0.0000001}) |
| self.assertEqual(ret["key"], "1e-07") |
| ret = TopologyType._sanitize_config({"key": -15.33333}) |
| self.assertEqual(ret["key"], "-15.33333") |
| |
| # non-string value -> should expect the same object |
| ret = TopologyType._sanitize_config({"key": ['v', 'a', 'l', 'u', 'e']}) |
| self.assertEqual(ret["key"], ['v', 'a', 'l', 'u', 'e']) |
| ret = TopologyType._sanitize_config({"key": None}) |
| self.assertEqual(ret["key"], None) |
| |
| def test_get_heron_options_from_env(self): |
| test_value = "cmdline.key.1=/tmp/directory,cmdline.with.space=hello%%%%world" |
| expecting = {"cmdline.key.1": "/tmp/directory", "cmdline.with.space": "hello world"} |
| os.environ["HERON_OPTIONS"] = test_value |
| ret = TopologyType.get_heron_options_from_env() |
| self.assertEqual(ret, expecting) |
| |
| # error |
| os.environ.pop("HERON_OPTIONS") |
| with self.assertRaises(RuntimeError): |
| TopologyType.get_heron_options_from_env() |
| |
| class TopologyBuilderTest(unittest.TestCase): |
| def test_constructor(self): |
| builder = TopologyBuilder("WordCount") |
| self.assertEqual(builder.topology_name, "WordCount") |
| |
| with self.assertRaises(AssertionError): |
| TopologyBuilder("Topology") |
| |
| with self.assertRaises(AssertionError): |
| TopologyBuilder(123) |
| |
| with self.assertRaises(AssertionError): |
| TopologyBuilder(None) |
| |
| def test_add_spec(self): |
| builder = TopologyBuilder("Test") |
| |
| with self.assertRaises(ValueError): |
| builder.add_spec(HeronComponentSpec(None, "path", True, 1)) |
| |
| with self.assertRaises(TypeError): |
| builder.add_spec(None) |
| |
| self.assertEqual(len(builder._specs), 0) |
| |
| # add 10 specs |
| specs = [] |
| for i in range(10): |
| specs.append(HeronComponentSpec(str(i), "path", True, 1)) |
| builder.add_spec(*specs) |
| self.assertEqual(len(builder._specs), 10) |
