beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From rober...@apache.org
Subject [3/3] beam git commit: [BEAM-1768] Fix assert_that for empty inputs
Date Tue, 21 Mar 2017 17:39:06 GMT
[BEAM-1768] Fix assert_that for empty inputs


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/c39b02d9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/c39b02d9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/c39b02d9

Branch: refs/heads/master
Commit: c39b02d9ea2d343f348beca1228776204605fb93
Parents: 030528f
Author: Robert Bradshaw <robertwb@gmail.com>
Authored: Mon Mar 20 18:20:01 2017 -0700
Committer: Robert Bradshaw <robertwb@google.com>
Committed: Tue Mar 21 10:38:37 2017 -0700

----------------------------------------------------------------------
 sdks/python/apache_beam/transforms/util.py      | 18 ++++---
 sdks/python/apache_beam/transforms/util_test.py | 50 ++++++++++++++++++++
 2 files changed, 62 insertions(+), 6 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/c39b02d9/sdks/python/apache_beam/transforms/util.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py
index e3f5b85..ac7eb3c 100644
--- a/sdks/python/apache_beam/transforms/util.py
+++ b/sdks/python/apache_beam/transforms/util.py
@@ -22,6 +22,7 @@ from __future__ import absolute_import
 
 from apache_beam.transforms import window
 from apache_beam.transforms.core import CombinePerKey
+from apache_beam.transforms.core import Create
 from apache_beam.transforms.core import Flatten
 from apache_beam.transforms.core import GroupByKey
 from apache_beam.transforms.core import Map
@@ -222,12 +223,17 @@ def assert_that(actual, matcher, label='assert_that'):
   class AssertThat(PTransform):
 
     def expand(self, pcoll):
-      return (pcoll
-              | WindowInto(window.GlobalWindows())
-              | "ToVoidKey" >> Map(lambda v: (None, v))
-              | "Group" >> GroupByKey()
-              | "UnKey" >> Map(lambda (k, v): v)
-              | "Match" >> Map(matcher))
+      # We must have at least a single element to ensure the matcher
+      # code gets run even if the input pcollection is empty.
+      keyed_singleton = pcoll.pipeline | Create([(None, None)])
+      keyed_actual = (
+          pcoll
+          | WindowInto(window.GlobalWindows())
+          | "ToVoidKey" >> Map(lambda v: (None, v)))
+      _ = ((keyed_singleton, keyed_actual)
+           | "Group" >> CoGroupByKey()
+           | "Unkey" >> Map(lambda (k, (_, actual_values)): actual_values)
+           | "Match" >> Map(matcher))
 
     def default_label(self):
       return label

http://git-wip-us.apache.org/repos/asf/beam/blob/c39b02d9/sdks/python/apache_beam/transforms/util_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py
new file mode 100644
index 0000000..9656827
--- /dev/null
+++ b/sdks/python/apache_beam/transforms/util_test.py
@@ -0,0 +1,50 @@
+#
+# 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.
+#
+
+"""Unit tests for the util transforms."""
+
+import unittest
+
+from apache_beam import Create
+from apache_beam.test_pipeline import TestPipeline
+from apache_beam.transforms.util import assert_that, equal_to, is_empty
+
+
+class UtilTest(unittest.TestCase):
+
+  def test_assert_that_passes(self):
+    with TestPipeline() as p:
+      assert_that(p | Create([1, 2, 3]), equal_to([1, 2, 3]))
+
+  def test_assert_that_fails(self):
+    with self.assertRaises(Exception):
+      with TestPipeline() as p:
+        assert_that(p | Create([1, 10, 100]), equal_to([1, 2, 3]))
+
+  def test_assert_that_fails_on_empty_input(self):
+    with self.assertRaises(Exception):
+      with TestPipeline() as p:
+        assert_that(p | Create([]), equal_to([1, 2, 3]))
+
+  def test_assert_that_fails_on_empty_expected(self):
+    with self.assertRaises(Exception):
+      with TestPipeline() as p:
+        assert_that(p | Create([1, 2, 3]), is_empty())
+
+
+if __name__ == '__main__':
+  unittest.main()


Mime
View raw message