From bc839499a557728408a9938b9fc16c70ee6df40b Mon Sep 17 00:00:00 2001
From: Wenjia Liu <wenjialiu@google.com>
Date: Fri, 15 Nov 2019 11:47:24 -0800
Subject: [PATCH] [BEAM-8575] Test DoFn context params

---
 sdks/python/apache_beam/pipeline_test.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py
index c1c25d1a6164e..a2d097c6a78a0 100644
--- a/sdks/python/apache_beam/pipeline_test.py
+++ b/sdks/python/apache_beam/pipeline_test.py
@@ -26,6 +26,7 @@
 from builtins import range
 
 import mock
+from nose.plugins.attrib import attr
 
 import apache_beam as beam
 from apache_beam import typehints
@@ -641,6 +642,26 @@ def process(self, element, prefix, suffix=DoFn.SideInputParam):
     assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
     pipeline.run()
 
+  @attr('ValidatesRunner')
+  def test_element_param(self):
+    pipeline = TestPipeline()
+    input = [1, 2]
+    pcoll = (pipeline
+             | 'Create' >> Create(input)
+             | 'Ele param' >> Map(lambda element=DoFn.ElementParam: element))
+    assert_that(pcoll, equal_to(input))
+    pipeline.run()
+
+  @attr('ValidatesRunner')
+  def test_key_param(self):
+    pipeline = TestPipeline()
+    pcoll = (pipeline
+             | 'Create' >> Create([('a', 1), ('b', 2)])
+             | 'Key param' >> Map(lambda _, key=DoFn.KeyParam: key))
+    assert_that(pcoll, equal_to(['a', 'b']))
+    pipeline.run()
+
+  @attr('ValidatesRunner')
   def test_window_param(self):
     class TestDoFn(DoFn):
       def process(self, element, window=DoFn.WindowParam):
@@ -663,6 +684,7 @@ def process(self, element, window=DoFn.WindowParam):
         label='doubled windows')
     pipeline.run()
 
+  @attr('ValidatesRunner')
   def test_timestamp_param(self):
     class TestDoFn(DoFn):
       def process(self, element, timestamp=DoFn.TimestampParam):
@@ -679,6 +701,7 @@ def test_timestamp_param_map(self):
           p | Create([1, 2]) | beam.Map(lambda _, t=DoFn.TimestampParam: t),
           equal_to([MIN_TIMESTAMP, MIN_TIMESTAMP]))
 
+  @attr('ValidatesRunner')
   def test_pane_info_param(self):
     with TestPipeline() as p:
       pc = p | Create([(None, None)])