add back tests

Author: Abacn

sdks/python/apache_beam/transforms/core.py

@@ -1520,9 +1520,8 @@ def _check_fn_use_yield_and_return(fn):
         if lstripped_line.rstrip() == "return None" or lstripped_line.rstrip(
         ) == "return(None)":
           has_return_none = True
-        else:
-          has_return = True
-      if has_yield and (has_return or has_return_none):
+        has_return = True
+      if has_yield and has_return:
         return True
 
     if has_return_none:

sdks/python/apache_beam/transforms/core_test.py

@@ -115,6 +115,12 @@ def process(self, element):
 
 
 class TestDoFn11(beam.DoFn):
+  """test process returning None (no return and no yield)"""
+  def process(self, element):
+    pass
+
+
+class TestDoFn12(beam.DoFn):
   """test process returning None in a filter pattern"""
   def process(self, element):
     if element == 0:
@@ -187,11 +193,16 @@ def test_dofn_with_explicit_return_none(self):
       assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text
       assert str(TestDoFn10) in self._caplog.text
 
-  def test_dofn_with_implicit_return_none_and_value(self):
+  def test_dofn_with_implicit_return_none_missing_return_and_yield(self):
     with self._caplog.at_level(logging.WARNING):
       beam.ParDo(TestDoFn11())
       assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text
 
+  def test_dofn_with_implicit_return_none_and_value(self):
+    with self._caplog.at_level(logging.WARNING):
+      beam.ParDo(TestDoFn12())
+      assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text
+
 
 class PartitionTest(unittest.TestCase):
   def test_partition_with_bools(self):