update a couple more tests

reeba212 · Dec 4, 2024 · 66b5320 · 66b5320
1 parent 4c9c5ce
commit 66b5320
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py b/sdks/python/apache_beam/typehints/typed_pipeline_test.py
@@ -422,19 +422,19 @@ def test_typed_ptransform_fn_conflicting_hints(self):
     # In this case, both MyMap and its contained ParDo have separate type
     # checks (that disagree with each other).
     @beam.ptransform_fn
-    @typehints.with_input_types(int)
+    @typehints.with_input_types(str)
     def MyMap(pcoll):
       def fn(element: float):
         yield element
 
       return pcoll | beam.ParDo(fn)
 
     with self.assertRaisesRegex(typehints.TypeCheckError,
-                                r'ParDo.*requires.*float.*got.*int'):
-      _ = [1, 2, 3] | MyMap()
+                                r'ParDo.*requires.*float.*got.*str'):
+      _ = ['1', '2', '3'] | MyMap()
     with self.assertRaisesRegex(typehints.TypeCheckError,
-                                r'MyMap.*expected.*int.*got.*str'):
-      _ = ['a'] | MyMap()
+                                r'MyMap.*expected.*str.*got.*bytes'):
+      _ = [b'a'] | MyMap()
 
   def test_typed_dofn_string_literals(self):
     class MyDoFn(beam.DoFn):

diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py
@@ -226,7 +226,7 @@ def test_union_hint_compatibility(self):
         typehints.Union[int, str],
         typehints.Union[str, typehints.Union[int, str]])
 
-    self.assertNotCompatible(
+    self.assertCompatible(
         typehints.Union[float, bool], typehints.Union[int, bool])
     self.assertNotCompatible(
         typehints.Union[bool, str], typehints.Union[float, bool, int])