You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by tv...@apache.org on 2022/02/01 22:38:56 UTC
[beam] branch master updated: [BEAM-12572] Run python examples on multiple runners (#16154)
This is an automated email from the ASF dual-hosted git repository.
tvalentyn 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 fa7e964 [BEAM-12572] Run python examples on multiple runners (#16154)
fa7e964 is described below
commit fa7e964b3d2fa2d215772d15f502e3b648ccdb7b
Author: Benjamin Gonzalez <74...@users.noreply.github.com>
AuthorDate: Tue Feb 1 16:37:37 2022 -0600
[BEAM-12572] Run python examples on multiple runners (#16154)
---
release/src/main/scripts/mass_comment.py | 4 ++
.../examples/complete/autocomplete_test.py | 80 ++++++++++++++++++----
.../apache_beam/examples/complete/estimate_pi.py | 2 +-
.../examples/complete/estimate_pi_test.py | 19 ++++-
.../examples/complete/game/game_stats_it_test.py | 6 ++
.../examples/complete/game/game_stats_test.py | 3 -
.../examples/complete/game/hourly_team_score.py | 4 +-
.../complete/game/hourly_team_score_it_test.py | 31 ++++++++-
.../complete/game/hourly_team_score_test.py | 3 -
.../examples/complete/game/leader_board_it_test.py | 6 ++
.../examples/complete/game/leader_board_test.py | 3 -
.../examples/complete/game/user_score.py | 4 +-
.../examples/complete/game/user_score_it_test.py | 29 +++++++-
.../examples/complete/game/user_score_test.py | 3 -
.../apache_beam/examples/complete/tfidf_test.py | 2 +-
.../complete/top_wikipedia_sessions_test.py | 31 ++++++++-
.../bigquery_side_input_it_test.py} | 52 ++++----------
.../examples/cookbook/bigquery_side_input_test.py | 3 -
.../examples/cookbook/bigquery_tornadoes_test.py | 3 -
.../apache_beam/examples/cookbook/coders_test.py | 50 +++++++++++++-
.../examples/cookbook/combiners_test.py | 3 -
.../examples/cookbook/custom_ptransform_test.py | 33 ++++++++-
.../apache_beam/examples/cookbook/filters_test.py | 43 +++++++++++-
.../examples/dataframe/taxiride_test.py | 2 -
.../examples/streaming_wordcount_debugging_test.py | 2 -
sdks/python/pytest.ini | 3 +
sdks/python/test-suites/direct/common.gradle | 4 +-
sdks/python/test-suites/portable/common.gradle | 4 +-
28 files changed, 335 insertions(+), 97 deletions(-)
diff --git a/release/src/main/scripts/mass_comment.py b/release/src/main/scripts/mass_comment.py
index bcfa099..866196b 100644
--- a/release/src/main/scripts/mass_comment.py
+++ b/release/src/main/scripts/mass_comment.py
@@ -59,6 +59,10 @@ COMMENTS_TO_ADD = [
"Run XVR_Flink PostCommit",
"Run XVR_Spark PostCommit",
"Run XVR_Spark3 PostCommit",
+ "Run Python Examples_Direct",
+ "Run Python Examples_Spark",
+ "Run Python Examples_Flink",
+ "Run Python Examples_Dataflow",
]
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete_test.py b/sdks/python/apache_beam/examples/complete/autocomplete_test.py
index d545edf..5ccbee7 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete_test.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete_test.py
@@ -19,6 +19,10 @@
# pytype: skip-file
+import logging
+import os
+import re
+import tempfile
import unittest
import pytest
@@ -29,6 +33,38 @@ from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_utils import compute_hash
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
+from apache_beam.testing.util import open_shards
+
+
+def format_output_file(output_string):
+ def extract_prefix_topk_words_tuples(line):
+ match = re.match(r'(.*): \[(.*)\]', line)
+ prefix = match.group(1)
+ topK_words_string = extract_top_k_words_tuples(match.group(2))
+ return prefix, topK_words_string
+
+ def extract_top_k_words_tuples(top_k_words_string):
+ top_k_list = top_k_words_string.split("), (")
+ return tuple(
+ map(
+ lambda top_k_string: tuple(format_top_k_tuples(top_k_string)),
+ top_k_list))
+
+ def format_top_k_tuples(top_k_string):
+ (frequency, words) = top_k_string.replace('(', '').replace(')', '').replace(
+ '\"', '').replace('\'', '').replace(' ', '').split(',')
+ return int(frequency), words
+
+ return list(
+ map(
+ lambda line: extract_prefix_topk_words_tuples(line),
+ output_string.split('\n')))
+
+
+def create_content_input_file(path, contents):
+ logging.info('Creating temp file: %s', path)
+ with open(path, 'w') as f:
+ f.write(contents)
class AutocompleteTest(unittest.TestCase):
@@ -36,25 +72,23 @@ class AutocompleteTest(unittest.TestCase):
WORDS = ['this', 'this', 'that', 'to', 'to', 'to']
KINGLEAR_HASH_SUM = 268011785062540
KINGLEAR_INPUT = 'gs://dataflow-samples/shakespeare/kinglear.txt'
+ EXPECTED_PREFIXES = [
+ ('t', ((3, 'to'), (2, 'this'), (1, 'that'))),
+ ('to', ((3, 'to'), )),
+ ('th', ((2, 'this'), (1, 'that'))),
+ ('thi', ((2, 'this'), )),
+ ('this', ((2, 'this'), )),
+ ('tha', ((1, 'that'), )),
+ ('that', ((1, 'that'), )),
+ ]
- @pytest.mark.examples_postcommit
def test_top_prefixes(self):
with TestPipeline() as p:
words = p | beam.Create(self.WORDS)
result = words | autocomplete.TopPerPrefix(5)
# values must be hashable for now
result = result | beam.Map(lambda k_vs: (k_vs[0], tuple(k_vs[1])))
- assert_that(
- result,
- equal_to([
- ('t', ((3, 'to'), (2, 'this'), (1, 'that'))),
- ('to', ((3, 'to'), )),
- ('th', ((2, 'this'), (1, 'that'))),
- ('thi', ((2, 'this'), )),
- ('this', ((2, 'this'), )),
- ('tha', ((1, 'that'), )),
- ('that', ((1, 'that'), )),
- ]))
+ assert_that(result, equal_to(self.EXPECTED_PREFIXES))
@pytest.mark.it_postcommit
def test_autocomplete_it(self):
@@ -71,6 +105,28 @@ class AutocompleteTest(unittest.TestCase):
assert_that(checksum, equal_to([self.KINGLEAR_HASH_SUM]))
+ @pytest.mark.examples_postcommit
+ def test_autocomplete_output_files_on_small_input(self):
+ logging.error('SAVE_MAIN_SESSION')
+ test_pipeline = TestPipeline(is_integration_test=True)
+ # Setup the files with expected content.
+ temp_folder = tempfile.mkdtemp()
+ create_content_input_file(
+ os.path.join(temp_folder, 'input.txt'), ' '.join(self.WORDS))
+ extra_opts = {
+ 'input': '%s/input.txt' % temp_folder,
+ 'output': os.path.join(temp_folder, 'result')
+ }
+
+ autocomplete.run(test_pipeline.get_full_options_as_args(**extra_opts))
+
+ # Load result file and compare.
+ with open_shards(os.path.join(temp_folder, 'result-*-of-*')) as result_file:
+ result = result_file.read().strip()
+
+ self.assertEqual(
+ sorted(self.EXPECTED_PREFIXES), sorted(format_output_file(result)))
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi.py b/sdks/python/apache_beam/examples/complete/estimate_pi.py
index 7f95989..089767d 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi.py
@@ -83,7 +83,7 @@ def combine_results(results):
class JsonCoder(object):
"""A JSON coder used to format the final result."""
def encode(self, x):
- return json.dumps(x)
+ return json.dumps(x).encode('utf-8')
class EstimatePiTransform(beam.PTransform):
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
index ac515c6..8b48542 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
@@ -19,7 +19,10 @@
# pytype: skip-file
+import json
import logging
+import os
+import tempfile
import unittest
import pytest
@@ -28,6 +31,7 @@ from apache_beam.examples.complete import estimate_pi
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import BeamAssertException
from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import open_shards
def in_between(lower, upper):
@@ -40,7 +44,6 @@ def in_between(lower, upper):
return _in_between
-@pytest.mark.examples_postcommit
class EstimatePiTest(unittest.TestCase):
def test_basics(self):
with TestPipeline() as p:
@@ -51,6 +54,20 @@ class EstimatePiTest(unittest.TestCase):
# trials.
assert_that(result, in_between(3.125, 3.155))
+ @pytest.mark.examples_postcommit
+ def test_estimate_pi_output_file(self):
+ test_pipeline = TestPipeline(is_integration_test=True)
+ temp_folder = tempfile.mkdtemp()
+ extra_opts = {'output': os.path.join(temp_folder, 'result')}
+ estimate_pi.run(test_pipeline.get_full_options_as_args(**extra_opts))
+ # Load result file and compare.
+ with open_shards(os.path.join(temp_folder, 'result-*-of-*')) as result_file:
+ [_, _, estimated_pi] = json.loads(result_file.read().strip())
+ # Note: Probabilistically speaking this test can fail with a probability
+ # that is very small (VERY) given that we run at least 100 thousand
+ # trials.
+ self.assertTrue(3.125 <= estimated_pi <= 3.155)
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
index b1b4746..c18eb40 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
@@ -105,6 +105,12 @@ class GameStatsIT(unittest.TestCase):
test_utils.cleanup_topics(self.pub_client, [self.input_topic])
@pytest.mark.it_postcommit
+ @pytest.mark.examples_postcommit
+ # TODO(BEAM-13613) This example only works in Dataflow,
+ # remove mark to enable for other runners when fixed
+ @pytest.mark.sickbay_direct
+ @pytest.mark.sickbay_spark
+ @pytest.mark.sickbay_flink
def test_game_stats_it(self):
state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats_test.py b/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
index ea2b971..d085b90 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.complete.game import game_stats
from apache_beam.testing.test_pipeline import TestPipeline
@@ -31,7 +29,6 @@ from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
-@pytest.mark.examples_postcommit
class GameStatsTest(unittest.TestCase):
SAMPLE_DATA = [
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
index 48a105a..31ab080 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
@@ -236,7 +236,7 @@ class HourlyTeamScore(beam.PTransform):
| 'ExtractAndSumScore' >> ExtractAndSumScore('team'))
-def run(argv=None, save_main_session=True):
+def run(argv=None):
"""Main entry point; defines and runs the hourly_team_score pipeline."""
parser = argparse.ArgumentParser()
@@ -293,7 +293,7 @@ def run(argv=None, save_main_session=True):
# We use the save_main_session option because one or more DoFn's in this
# workflow rely on global context (e.g., a module imported at module level).
- options.view_as(SetupOptions).save_main_session = save_main_session
+ options.view_as(SetupOptions).save_main_session = True
with beam.Pipeline(options=options) as p:
( # pylint: disable=expression-not-assigned
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
index a9729bc..56ef4a5 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
@@ -87,8 +87,35 @@ class HourlyTeamScoreIT(unittest.TestCase):
# Get pipeline options from command argument: --test-pipeline-options,
# and start pipeline job by calling pipeline main function.
hourly_team_score.run(
- self.test_pipeline.get_full_options_as_args(**extra_opts),
- save_main_session=False)
+ self.test_pipeline.get_full_options_as_args(**extra_opts))
+
+ @pytest.mark.examples_postcommit
+ def test_hourly_team_score_output_checksum_on_small_input(self):
+ # Small dataset to prevent Out of Memory when running in local runners
+ INPUT_FILE = 'gs://apache-beam-samples/game/small/gaming_data.csv'
+ EXPECTED_CHECKSUM = '91143e81622aa391eb62eaa3f3a5123401edb07d'
+ state_verifier = PipelineStateMatcher(PipelineState.DONE)
+ query = (
+ 'SELECT COUNT(*) FROM `%s.%s.%s`' %
+ (self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE))
+
+ bigquery_verifier = BigqueryMatcher(self.project, query, EXPECTED_CHECKSUM)
+
+ extra_opts = {
+ 'input': INPUT_FILE,
+ 'dataset': self.dataset_ref.dataset_id,
+ 'window_duration': 1,
+ 'on_success_matcher': all_of(state_verifier, bigquery_verifier)
+ }
+
+ # Register clean up before pipeline execution
+ # Note that actual execution happens in reverse order.
+ self.addCleanup(utils.delete_bq_dataset, self.project, self.dataset_ref)
+
+ # Get pipeline options from command argument: --test-pipeline-options,
+ # and start pipeline job by calling pipeline main function.
+ hourly_team_score.run(
+ self.test_pipeline.get_full_options_as_args(**extra_opts))
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
index 55028f9..804cc84 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.complete.game import hourly_team_score
from apache_beam.testing.test_pipeline import TestPipeline
@@ -42,7 +40,6 @@ class HourlyTeamScoreTest(unittest.TestCase):
'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955',
]
- @pytest.mark.examples_postcommit
def test_hourly_team_score(self):
with TestPipeline() as p:
result = (
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
index 3c1677c..df686fc 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
@@ -106,6 +106,12 @@ class LeaderBoardIT(unittest.TestCase):
test_utils.cleanup_topics(self.pub_client, [self.input_topic])
@pytest.mark.it_postcommit
+ @pytest.mark.examples_postcommit
+ # TODO(BEAM-13613) This example only works in Dataflow,
+ # remove mark to enable for other runners when fixed
+ @pytest.mark.sickbay_direct
+ @pytest.mark.sickbay_spark
+ @pytest.mark.sickbay_flink
def test_leader_board_it(self):
state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board_test.py b/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
index 4a6c44b..1c1cd65 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.complete.game import leader_board
from apache_beam.options.pipeline_options import PipelineOptions
@@ -32,7 +30,6 @@ from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
-@pytest.mark.examples_postcommit
class LeaderBoardTest(unittest.TestCase):
SAMPLE_DATA = [
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score.py b/sdks/python/apache_beam/examples/complete/game/user_score.py
index 6a97d7e..b645809 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score.py
@@ -130,7 +130,7 @@ class UserScore(beam.PTransform):
# [START main]
-def run(argv=None, save_main_session=True):
+def run(argv=None):
"""Main entry point; defines and runs the user_score pipeline."""
parser = argparse.ArgumentParser()
@@ -150,7 +150,7 @@ def run(argv=None, save_main_session=True):
# We use the save_main_session option because one or more DoFn's in this
# workflow rely on global context (e.g., a module imported at module level).
- options.view_as(SetupOptions).save_main_session = save_main_session
+ options.view_as(SetupOptions).save_main_session = True
with beam.Pipeline(options=options) as p:
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py b/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
index d26565e..75fceb5 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
@@ -51,14 +51,15 @@ from apache_beam.testing.test_utils import delete_files
class UserScoreIT(unittest.TestCase):
DEFAULT_INPUT_FILE = 'gs://dataflow-samples/game/gaming_data*'
+ DEFAULT_OUTPUT_FILE = \
+ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output'
DEFAULT_EXPECTED_CHECKSUM = '9f3bd81669607f0d98ec80ddd477f3277cfba0a2'
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.uuid = str(uuid.uuid4())
- self.output = '/'.join(
- [self.test_pipeline.get_option('output'), self.uuid, 'results'])
+ self.output = '/'.join([self.DEFAULT_OUTPUT_FILE, self.uuid, 'results'])
@pytest.mark.it_postcommit
def test_user_score_it(self):
@@ -84,6 +85,30 @@ class UserScoreIT(unittest.TestCase):
self.test_pipeline.get_full_options_as_args(**extra_opts),
save_main_session=False)
+ @pytest.mark.examples_postcommit
+ def test_userscore_output_checksum_on_small_input(self):
+ # Small dataset to prevent Out of Memory when running in local runners
+ INPUT_FILE = 'gs://apache-beam-samples/game/small/gaming_data.csv'
+ EXPECTED_CHECKSUM = '5b1bc0e8080e3c0f162809ac4c0f49acab23854e'
+ state_verifier = PipelineStateMatcher(PipelineState.DONE)
+ arg_sleep_secs = self.test_pipeline.get_option('sleep_secs')
+ sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None
+ file_verifier = FileChecksumMatcher(
+ self.output + '/*-of-*', EXPECTED_CHECKSUM, sleep_secs)
+
+ extra_opts = {
+ 'input': INPUT_FILE,
+ 'output': self.output + '/user-score',
+ 'on_success_matcher': all_of(state_verifier, file_verifier)
+ }
+
+ # Register clean up before pipeline execution
+ self.addCleanup(delete_files, [self.output + '*'])
+
+ # Get pipeline options from command argument: --test-pipeline-options,
+ # and start pipeline job by calling pipeline main function.
+ user_score.run(self.test_pipeline.get_full_options_as_args(**extra_opts))
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.DEBUG)
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score_test.py b/sdks/python/apache_beam/examples/complete/game/user_score_test.py
index b96f2e9..2283e56 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.complete.game import user_score
from apache_beam.testing.test_pipeline import TestPipeline
@@ -42,7 +40,6 @@ class UserScoreTest(unittest.TestCase):
'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955',
]
- @pytest.mark.examples_postcommit
def test_user_score(self):
with TestPipeline() as p:
result = (
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_test.py b/sdks/python/apache_beam/examples/complete/tfidf_test.py
index a55b076..07138cd 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf_test.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf_test.py
@@ -43,7 +43,6 @@ EXPECTED_RESULTS = set([
EXPECTED_LINE_RE = r'\(u?\'([a-z]*)\', \(\'.*([0-9]\.txt)\', (.*)\)\)'
-@pytest.mark.examples_postcommit
class TfIdfTest(unittest.TestCase):
def create_file(self, path, contents):
logging.info('Creating temp file: %s', path)
@@ -66,6 +65,7 @@ class TfIdfTest(unittest.TestCase):
# To actually trigger the check the pipeline must be run (e.g. by
# exiting the with context).
+ @pytest.mark.examples_postcommit
def test_basics(self):
# Setup the files with expected content.
temp_folder = tempfile.mkdtemp()
diff --git a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
index 0e27db6..9a4e5f4 100644
--- a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
+++ b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
@@ -20,6 +20,9 @@
# pytype: skip-file
import json
+import logging
+import os
+import tempfile
import unittest
import pytest
@@ -29,6 +32,7 @@ from apache_beam.examples.complete import top_wikipedia_sessions
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
+from apache_beam.testing.util import open_shards
class ComputeTopSessionsTest(unittest.TestCase):
@@ -70,7 +74,11 @@ class ComputeTopSessionsTest(unittest.TestCase):
'user3 : [3024.0, 6624.0) : 1 : [0.0, 2592000.0)',
]
- @pytest.mark.examples_postcommit
+ def create_content_input_file(self, path, contents):
+ logging.info('Creating temp file: %s', path)
+ with open(path, 'w') as f:
+ f.write(contents)
+
def test_compute_top_sessions(self):
with TestPipeline() as p:
edits = p | beam.Create(self.EDITS)
@@ -78,6 +86,27 @@ class ComputeTopSessionsTest(unittest.TestCase):
assert_that(result, equal_to(self.EXPECTED))
+ @pytest.mark.examples_postcommit
+ def test_top_wikipedia_sessions_output_files_on_small_input(self):
+ test_pipeline = TestPipeline(is_integration_test=True)
+ # Setup the files with expected content.
+ temp_folder = tempfile.mkdtemp()
+ self.create_content_input_file(
+ os.path.join(temp_folder, 'input.txt'), '\n'.join(self.EDITS))
+ extra_opts = {
+ 'input': '%s/input.txt' % temp_folder,
+ 'output': os.path.join(temp_folder, 'result'),
+ 'sampling_threshold': '1.0'
+ }
+ top_wikipedia_sessions.run(
+ test_pipeline.get_full_options_as_args(**extra_opts))
+
+ # Load result file and compare.
+ with open_shards(os.path.join(temp_folder, 'result-*-of-*')) as result_file:
+ result = result_file.read().strip().splitlines()
+
+ self.assertEqual(self.EXPECTED, sorted(result, key=lambda x: x.split()[0]))
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_it_test.py
similarity index 55%
copy from sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
copy to sdks/python/apache_beam/examples/cookbook/bigquery_side_input_it_test.py
index d26565e..c326cd5 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_it_test.py
@@ -15,21 +15,7 @@
# limitations under the License.
#
-"""End-to-end test for the user score example.
-
-Code: beam/sdks/python/apache_beam/examples/complete/game/user_score.py
-Usage:
-
- pytest --test-pipeline-options=" \
- --runner=TestDataflowRunner \
- --project=... \
- --region=... \
- --staging_location=gs://... \
- --temp_location=gs://... \
- --output=gs://... \
- --sdk_location=... \
-
-"""
+"""Integration test for the BigQuery side input example."""
# pytype: skip-file
@@ -40,39 +26,32 @@ import uuid
import pytest
from hamcrest.core.core.allof import all_of
-from apache_beam.examples.complete.game import user_score
+from apache_beam.examples.cookbook import bigquery_side_input
from apache_beam.runners.runner import PipelineState
-from apache_beam.testing.pipeline_verifiers import FileChecksumMatcher
from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_utils import delete_files
-class UserScoreIT(unittest.TestCase):
-
- DEFAULT_INPUT_FILE = 'gs://dataflow-samples/game/gaming_data*'
- DEFAULT_EXPECTED_CHECKSUM = '9f3bd81669607f0d98ec80ddd477f3277cfba0a2'
+class BigQuerySideInputIT(unittest.TestCase):
+ DEFAULT_OUTPUT_FILE = \
+ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output'
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.uuid = str(uuid.uuid4())
- self.output = '/'.join(
- [self.test_pipeline.get_option('output'), self.uuid, 'results'])
-
- @pytest.mark.it_postcommit
- def test_user_score_it(self):
+ self.output = '/'.join([self.DEFAULT_OUTPUT_FILE, self.uuid, 'results'])
+ @pytest.mark.examples_postcommit
+ def test_bigquery_side_input_it(self):
state_verifier = PipelineStateMatcher(PipelineState.DONE)
- arg_sleep_secs = self.test_pipeline.get_option('sleep_secs')
- sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None
- file_verifier = FileChecksumMatcher(
- self.output + '/*-of-*', self.DEFAULT_EXPECTED_CHECKSUM, sleep_secs)
+ NUM_GROUPS = 3
extra_opts = {
- 'input': self.DEFAULT_INPUT_FILE,
- 'output': self.output + '/user-score',
- 'on_success_matcher': all_of(state_verifier, file_verifier)
+ 'output': self.output,
+ 'num_groups': str(NUM_GROUPS),
+ 'on_success_matcher': all_of(state_verifier)
}
# Register clean up before pipeline execution
@@ -80,11 +59,10 @@ class UserScoreIT(unittest.TestCase):
# Get pipeline options from command argument: --test-pipeline-options,
# and start pipeline job by calling pipeline main function.
- user_score.run(
- self.test_pipeline.get_full_options_as_args(**extra_opts),
- save_main_session=False)
+ bigquery_side_input.run(
+ self.test_pipeline.get_full_options_as_args(**extra_opts))
if __name__ == '__main__':
- logging.getLogger().setLevel(logging.DEBUG)
+ logging.getLogger().setLevel(logging.INFO)
unittest.main()
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
index 61b7bf9..ba9b61e 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.cookbook import bigquery_side_input
from apache_beam.testing.test_pipeline import TestPipeline
@@ -32,7 +30,6 @@ from apache_beam.testing.util import equal_to
class BigQuerySideInputTest(unittest.TestCase):
- @pytest.mark.examples_postcommit
def test_create_groups(self):
with TestPipeline() as p:
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
index 6ca98e9..e9f3f67 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
@@ -22,8 +22,6 @@
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.examples.cookbook import bigquery_tornadoes
from apache_beam.testing.test_pipeline import TestPipeline
@@ -32,7 +30,6 @@ from apache_beam.testing.util import equal_to
class BigQueryTornadoesTest(unittest.TestCase):
- @pytest.mark.examples_postcommit
def test_basics(self):
with TestPipeline() as p:
rows = (
diff --git a/sdks/python/apache_beam/examples/cookbook/coders_test.py b/sdks/python/apache_beam/examples/cookbook/coders_test.py
index 5bbfd5e..0911218 100644
--- a/sdks/python/apache_beam/examples/cookbook/coders_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/coders_test.py
@@ -19,7 +19,10 @@
# pytype: skip-file
+import json
import logging
+import os
+import tempfile
import unittest
import pytest
@@ -29,6 +32,7 @@ from apache_beam.examples.cookbook import coders
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
+from apache_beam.testing.util import open_shards
class CodersTest(unittest.TestCase):
@@ -41,7 +45,26 @@ class CodersTest(unittest.TestCase):
'host': ['Brasil', 1], 'guest': ['Italy', 0]
}]
- @pytest.mark.examples_postcommit
+ EXPECTED_RESULT = [('Italy', 0), ('Brasil', 6), ('Germany', 3)]
+
+ def create_content_input_file(self, path, contents):
+ logging.info('Creating temp file: %s', path)
+ with open(path, 'w') as f:
+ f.write(contents)
+
+ def format_result(self, result_string):
+ def format_tuple(result_elem_list):
+ [country, counter] = result_elem_list
+ return country, int(counter.strip())
+
+ result_list = list(
+ map(
+ lambda result_elem: format_tuple(result_elem.split(',')),
+ result_string.replace('\'',
+ '').replace('[', '').replace(']', '').replace(
+ '\"', '').split('\n')))
+ return result_list
+
def test_compute_points(self):
with TestPipeline() as p:
records = p | 'create' >> beam.Create(self.SAMPLE_RECORDS)
@@ -49,8 +72,29 @@ class CodersTest(unittest.TestCase):
records
| 'points' >> beam.FlatMap(coders.compute_points)
| beam.CombinePerKey(sum))
- assert_that(
- result, equal_to([('Italy', 0), ('Brasil', 6), ('Germany', 3)]))
+ assert_that(result, equal_to(self.EXPECTED_RESULT))
+
+ @pytest.mark.examples_postcommit
+ def test_coders_output_files_on_small_input(self):
+ test_pipeline = TestPipeline(is_integration_test=True)
+
+ # Setup the files with expected content.
+ temp_folder = tempfile.mkdtemp()
+ self.create_content_input_file(
+ os.path.join(temp_folder, 'input.txt'),
+ '\n'.join(map(json.dumps, self.SAMPLE_RECORDS)))
+ extra_opts = {
+ 'input': '%s/input.txt' % temp_folder,
+ 'output': os.path.join(temp_folder, 'result')
+ }
+ coders.run(test_pipeline.get_full_options_as_args(**extra_opts))
+
+ # Load result file and compare.
+ with open_shards(os.path.join(temp_folder, 'result-*-of-*')) as result_file:
+ result = result_file.read().strip()
+
+ self.assertEqual(
+ sorted(self.EXPECTED_RESULT), sorted(self.format_result(result)))
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/cookbook/combiners_test.py b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
index e5c4dde..e844635 100644
--- a/sdks/python/apache_beam/examples/cookbook/combiners_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
@@ -35,15 +35,12 @@ checked directly on the last PCollection produced.
import logging
import unittest
-import pytest
-
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
-@pytest.mark.examples_postcommit
class CombinersTest(unittest.TestCase):
"""Tests showcasing Dataflow combiners."""
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
index a5c0d4d..4bcd095 100644
--- a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
@@ -20,6 +20,8 @@
# pytype: skip-file
import logging
+import os
+import tempfile
import unittest
import pytest
@@ -29,10 +31,17 @@ from apache_beam.examples.cookbook import custom_ptransform
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
+from apache_beam.testing.util import open_shards
-@pytest.mark.examples_postcommit
class CustomCountTest(unittest.TestCase):
+ WORDS = ['CAT', 'DOG', 'CAT', 'CAT', 'DOG']
+
+ def create_content_input_file(self, path, contents):
+ logging.info('Creating temp file: %s', path)
+ with open(path, 'w') as f:
+ f.write(contents)
+
def test_count1(self):
self.run_pipeline(custom_ptransform.Count1())
@@ -45,11 +54,31 @@ class CustomCountTest(unittest.TestCase):
def run_pipeline(self, count_implementation, factor=1):
with TestPipeline() as p:
- words = p | beam.Create(['CAT', 'DOG', 'CAT', 'CAT', 'DOG'])
+ words = p | beam.Create(self.WORDS)
result = words | count_implementation
assert_that(
result, equal_to([('CAT', (3 * factor)), ('DOG', (2 * factor))]))
+ @pytest.mark.examples_postcommit
+ def test_custom_ptransform_output_files_on_small_input(self):
+ EXPECTED_RESULT = "('CAT DOG CAT CAT DOG', 2)"
+
+ # Setup the files with expected content.
+ temp_folder = tempfile.mkdtemp()
+ self.create_content_input_file(
+ os.path.join(temp_folder, 'input.txt'), ' '.join(self.WORDS))
+ custom_ptransform.run([
+ '--input=%s/input.txt' % temp_folder,
+ '--output',
+ os.path.join(temp_folder, 'result')
+ ])
+
+ # Load result file and compare.
+ with open_shards(os.path.join(temp_folder, 'result-*-of-*')) as result_file:
+ result = result_file.read().strip()
+
+ self.assertEqual(result, EXPECTED_RESULT)
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/examples/cookbook/filters_test.py b/sdks/python/apache_beam/examples/cookbook/filters_test.py
index 0a0bbb4..f76afd8 100644
--- a/sdks/python/apache_beam/examples/cookbook/filters_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/filters_test.py
@@ -27,19 +27,28 @@
# - Filtering
import logging
+import time
import unittest
import pytest
+from hamcrest.core.core.allof import all_of
import apache_beam as beam
from apache_beam.examples.cookbook import filters
+from apache_beam.io.gcp.tests import utils
+from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher
+from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
-@pytest.mark.examples_postcommit
class FiltersTest(unittest.TestCase):
+
+ # The default checksum is a SHA-1 hash generated from sorted rows reading
+ # from expected Bigquery table.
+ DEFAULT_CHECKSUM = '813b6da1624334732fad4467d74a7c8a62559c6b'
+
# Note that 'removed' should be projected away by the pipeline
input_data = [
{
@@ -64,7 +73,7 @@ class FiltersTest(unittest.TestCase):
results = filters.filter_cold_days(rows, month)
return results
- def test_basic(self):
+ def test_basics(self):
"""Test that the correct result is returned for a simple dataset."""
with TestPipeline() as p:
results = self._get_result_for_month(p, 1)
@@ -88,6 +97,36 @@ class FiltersTest(unittest.TestCase):
results = self._get_result_for_month(p, 4)
assert_that(results, equal_to([]))
+ @pytest.mark.examples_postcommit
+ def test_filters_output_bigquery_matcher(self):
+ test_pipeline = TestPipeline(is_integration_test=True)
+
+ # Set extra options to the pipeline for test purpose
+ project = test_pipeline.get_option('project')
+
+ dataset = 'FiltersTestIT'
+ table = 'cold_days_%s' % int(round(time.time() * 1000))
+ output_table = '.'.join([dataset, table])
+ query = 'SELECT year, month, day, mean_temp FROM `%s`' % output_table
+
+ pipeline_verifiers = [
+ PipelineStateMatcher(),
+ BigqueryMatcher(
+ project=project, query=query, checksum=self.DEFAULT_CHECKSUM)
+ ]
+ extra_opts = {
+ 'output': output_table,
+ 'on_success_matcher': all_of(*pipeline_verifiers)
+ }
+
+ # Register cleanup before pipeline execution.
+ # Note that actual execution happens in reverse order.
+ self.addCleanup(utils.delete_bq_table, project, dataset, table)
+
+ # Get pipeline options from command argument: --test-pipeline-options,
+ # and start pipeline job by calling pipeline main function.
+ filters.run(test_pipeline.get_full_options_as_args(**extra_opts))
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/examples/dataframe/taxiride_test.py b/sdks/python/apache_beam/examples/dataframe/taxiride_test.py
index 1df2ddb..5334834 100644
--- a/sdks/python/apache_beam/examples/dataframe/taxiride_test.py
+++ b/sdks/python/apache_beam/examples/dataframe/taxiride_test.py
@@ -30,14 +30,12 @@ import tempfile
import unittest
import pandas as pd
-import pytest
import apache_beam as beam
from apache_beam.examples.dataframe import taxiride
from apache_beam.testing.util import open_shards
-@pytest.mark.examples_postcommit
class TaxiRideExampleTest(unittest.TestCase):
# First 10 lines from gs://apache-beam-samples/nyc_taxi/misc/sample.csv
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
index a10e01e..2ab8099 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
@@ -22,7 +22,6 @@
import unittest
import mock
-import pytest
import apache_beam as beam
from apache_beam.examples import streaming_wordcount_debugging
@@ -39,7 +38,6 @@ except ImportError:
# pylint: enable=wrong-import-order, wrong-import-position
-@pytest.mark.examples_postcommit
class StreamingWordcountDebugging(unittest.TestCase):
@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
@mock.patch('apache_beam.io.ReadFromPubSub')
diff --git a/sdks/python/pytest.ini b/sdks/python/pytest.ini
index ec32976..36e09d5 100644
--- a/sdks/python/pytest.ini
+++ b/sdks/python/pytest.ini
@@ -36,6 +36,9 @@ markers =
no_sickbay_batch: run without sickbay-batch
it_validatescontainer: collect for ValidatesContainer integration test runs
examples_postcommit: collect for post-commit test examples runs
+ sickbay_direct: run without sickbay-direct
+ sickbay_spark: run without sickbay-spark
+ sickbay_flink: run without sickbay-flink
# Tests using this marker conflict with the xdist plugin in some way, such
# as enabling save_main_session.
no_xdist: run without pytest-xdist plugin
diff --git a/sdks/python/test-suites/direct/common.gradle b/sdks/python/test-suites/direct/common.gradle
index 3ef5451..d90cd03 100644
--- a/sdks/python/test-suites/direct/common.gradle
+++ b/sdks/python/test-suites/direct/common.gradle
@@ -82,11 +82,11 @@ task examples {
dependsOn ':sdks:python:sdist'
doLast {
- def testOpts = basicTestOpts
+ def testOpts = ["--log-cli-level=INFO"]
def argMap = ["runner": "TestDirectRunner",
"test_opts": testOpts,
"suite": "postCommitExamples-direct-py${pythonVersionSuffix}",
- "collect": "examples_postcommit"
+ "collect": "examples_postcommit and not sickbay_direct"
]
def batchCmdArgs = mapToArgString(argMap)
exec {
diff --git a/sdks/python/test-suites/portable/common.gradle b/sdks/python/test-suites/portable/common.gradle
index be9feaf..32e566a 100644
--- a/sdks/python/test-suites/portable/common.gradle
+++ b/sdks/python/test-suites/portable/common.gradle
@@ -206,7 +206,7 @@ project.task("flinkExamples") {
"test_opts": testOpts,
"suite": "postCommitExamples-flink-py${pythonVersionSuffix}",
"pipeline_opts": pipelineOpts.join(" "),
- "collect": "examples_postcommit"
+ "collect": "examples_postcommit and not sickbay_flink"
])
exec {
executable 'sh'
@@ -237,7 +237,7 @@ project.task("sparkExamples") {
"test_opts": testOpts,
"suite": "postCommitExamples-spark-py${pythonVersionSuffix}",
"pipeline_opts": pipelineOpts.join(" "),
- "collect": "examples_postcommit"
+ "collect": "examples_postcommit and not sickbay_spark"
])
exec {
executable 'sh'