Cleaned the integration

Osmar Coronel · Osmar Coronel · commit 80260b41bc04 · 2019-08-09T01:29:36.000-07:00
diff --git a/.travis.yml b/.travis.yml
@@ -18,7 +18,6 @@ branches:
   only:
     - master
     - /^release\/.+$/
-    - beam-integration
 
 matrix:
   allow_failures:
diff --git a/sentry_sdk/integrations/beam.py b/sentry_sdk/integrations/beam.py
@@ -30,7 +30,7 @@ def setup_once():
             setattr(
                 DoFn,
                 INSPECT_FUNC.format(func_name),
-                patched_inspect_process(DoFn, func_name),
+                _wrap_inspect_call(DoFn, func_name),
             )
 
         old_init = ParDo.__init__
@@ -61,7 +61,7 @@ def sentry_init_pardo(self, fn, *args, **kwargs):
         ParDo.__init__ = sentry_init_pardo
 
 
-def patched_inspect_process(cls, func_name):
+def _wrap_inspect_call(cls, func_name):
     from apache_beam.typehints.decorators import getfullargspec  # type: ignore
 
     if not hasattr(cls, func_name):
@@ -108,7 +108,7 @@ def _inner(*args, **kwargs):
 
 def _capture_exception(exc_info, hub):
     """
-    Send Beam exception to Sentry
+    Send Beam exception to Sentry.
     """
     integration = hub.get_integration(BeamIntegration)
     if integration:
@@ -142,6 +142,6 @@ def _wrap_generator_call(gen, client):
         try:
             yield next(gen)
         except StopIteration:
-            raise
+            break
         except Exception:
             raise_exception(client)
diff --git a/tests/integrations/beam/test_beam.py b/tests/integrations/beam/test_beam.py
@@ -1,13 +1,14 @@
 import pytest
-import dill
 import inspect
 
 pytest.importorskip("apache_beam")
 
+import dill
+
 from sentry_sdk.integrations.beam import (
     BeamIntegration,
     _wrap_task_call,
-    patched_inspect_process,
+    _wrap_inspect_call,
 )
 
 from apache_beam.typehints.trivial_inference import instance_to_type
@@ -35,7 +36,7 @@ class A:
     def __init__(self, fn):
         self.r = "We are in A"
         self.fn = fn
-        setattr(self, "_inspect_fn", patched_inspect_process(self, "fn"))
+        setattr(self, "_inspect_fn", _wrap_inspect_call(self, "fn"))
 
     def process(self):
         return self.fn()
@@ -153,8 +154,12 @@ def test_monkey_patch_signature(f, args, kwargs):
 class _OutputProcessor(OutputProcessor):
     def process_outputs(self, windowed_input_element, results):
         print(windowed_input_element)
-        for result in results:
-            assert result
+        try:
+            for result in results:
+                assert result
+        except StopIteration:
+            print("In here")
+
 
 
 @pytest.fixture
diff --git a/tox.ini b/tox.ini
@@ -32,8 +32,8 @@ envlist =
     {pypy,py2.7,py3.5,py3.6,py3.7,py3.8}-celery-{4.1,4.2,4.3}
     {pypy,py2.7}-celery-3
 
-    py2.7-beam-{12,13,master}
-    {py3.6,py3.7,py3.8}-beam-{12,13,master,dev}
+    {py2.7,py3.6}-beam-{12,13,master}
+    py3.7-beam-{12,13}
 
     # The aws_lambda tests deploy to the real AWS and have their own matrix of Python versions.
     py3.7-aws_lambda
