[beam] branch master updated: Conditionally run Python PubSub protobuf unit test

[pabloem]

This is an automated email from the ASF dual-hosted git repository.

pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 9e07699  Conditionally run Python PubSub protobuf unit test
9e07699 is described below

commit 9e0769936bbde7a05becc518f613e4c53634e2c0
Author: Udi Meiri <eh...@google.com>
AuthorDate: Tue Aug 7 18:42:56 2018 -0700

    Conditionally run Python PubSub protobuf unit test
    
    Skips test if protobuf packages are not installed.
---
 sdks/python/apache_beam/io/gcp/pubsub.py      | 1 +
 sdks/python/apache_beam/io/gcp/pubsub_test.py | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py 
b/sdks/python/apache_beam/io/gcp/pubsub.py
index 3b65fac..d5c3f99 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -39,6 +39,7 @@ from apache_beam.transforms import PTransform
 from apache_beam.transforms.display import DisplayDataItem
 from apache_beam.utils.annotations import deprecated
 
+# The protobuf library is only used for running on Dataflow.
 try:
   from google.cloud.proto.pubsub.v1 import pubsub_pb2
 except ImportError:
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py 
b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index c1dd05e..914b72a 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -61,6 +61,12 @@ except ImportError:
   pubsub = None
 # pylint: enable=wrong-import-order, wrong-import-position
 
+# The protobuf library is only used for running on Dataflow.
+try:
+  from google.cloud.proto.pubsub.v1 import pubsub_pb2
+except ImportError:
+  pubsub_pb2 = None
+
 
 class TestPubsubMessage(unittest.TestCase):
 
@@ -75,7 +81,8 @@ class TestPubsubMessage(unittest.TestCase):
     with self.assertRaisesRegexp(ValueError, r'data.*attributes.*must be set'):
       _ = PubsubMessage(None, {})
 
-  @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
+  @unittest.skipIf(pubsub_pb2 is None,
+                   'PubSub proto dependencies are not installed')
   def test_proto_conversion(self):
     data = 'data'
     attributes = {'k1': 'v1', 'k2': 'v2'}