Cleanup temp file handling in tests. Also reduces flakiness when temp files were not properly isolated.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/151bc263 Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/151bc263 Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/151bc263 Branch: refs/heads/master Commit: 151bc26367cf6d88916b7c67f6c5df742f8440c2 Parents: 603f86a Author: Robert Bradshaw <rober...@gmail.com> Authored: Tue Nov 21 10:50:04 2017 -0800 Committer: Robert Bradshaw <rober...@gmail.com> Committed: Tue Nov 21 16:07:49 2017 -0800 ---------------------------------------------------------------------- .../dataflow/internal/dependency_test.py | 108 ++++++++++--------- 1 file changed, 57 insertions(+), 51 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/151bc263/sdks/python/apache_beam/runners/dataflow/internal/dependency_test.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/runners/dataflow/internal/dependency_test.py b/sdks/python/apache_beam/runners/dataflow/internal/dependency_test.py index f0e59bc..68e5d8c 100644 --- a/sdks/python/apache_beam/runners/dataflow/internal/dependency_test.py +++ b/sdks/python/apache_beam/runners/dataflow/internal/dependency_test.py @@ -42,6 +42,18 @@ except ImportError: @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class SetupTest(unittest.TestCase): + def setUp(self): + self._temp_dir = None + + def make_temp_dir(self): + if self._temp_dir is None: + self._temp_dir = tempfile.mkdtemp() + return tempfile.mkdtemp(dir=self._temp_dir) + + def tearDown(self): + if self._temp_dir: + shutil.rmtree(self._temp_dir) + def update_options(self, options): setup_options = options.view_as(SetupOptions) setup_options.sdk_location = '' @@ -66,7 +78,7 @@ class SetupTest(unittest.TestCase): cm.exception.message) def test_no_temp_location(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() options = PipelineOptions() google_cloud_options = options.view_as(GoogleCloudOptions) google_cloud_options.staging_location = staging_dir @@ -78,7 +90,7 @@ class SetupTest(unittest.TestCase): cm.exception.message) def test_no_main_session(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir @@ -90,7 +102,7 @@ class SetupTest(unittest.TestCase): dependency.stage_job_resources(options)) def test_with_main_session(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir @@ -105,7 +117,7 @@ class SetupTest(unittest.TestCase): os.path.join(staging_dir, names.PICKLED_MAIN_SESSION_FILE))) def test_default_resources(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) @@ -115,37 +127,32 @@ class SetupTest(unittest.TestCase): dependency.stage_job_resources(options)) def test_with_requirements_file(self): - try: - staging_dir = tempfile.mkdtemp() - requirements_cache_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + requirements_cache_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() - options = PipelineOptions() - options.view_as(GoogleCloudOptions).staging_location = staging_dir - self.update_options(options) - options.view_as(SetupOptions).requirements_cache = requirements_cache_dir - options.view_as(SetupOptions).requirements_file = os.path.join( - source_dir, dependency.REQUIREMENTS_FILE) - self.create_temp_file( - os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing') - self.assertEqual( - sorted([dependency.REQUIREMENTS_FILE, - 'abc.txt', 'def.txt']), - sorted(dependency.stage_job_resources( - options, - populate_requirements_cache=self.populate_requirements_cache))) - self.assertTrue( - os.path.isfile( - os.path.join(staging_dir, dependency.REQUIREMENTS_FILE))) - self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) - self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) - finally: - shutil.rmtree(staging_dir) - shutil.rmtree(requirements_cache_dir) - shutil.rmtree(source_dir) + options = PipelineOptions() + options.view_as(GoogleCloudOptions).staging_location = staging_dir + self.update_options(options) + options.view_as(SetupOptions).requirements_cache = requirements_cache_dir + options.view_as(SetupOptions).requirements_file = os.path.join( + source_dir, dependency.REQUIREMENTS_FILE) + self.create_temp_file( + os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing') + self.assertEqual( + sorted([dependency.REQUIREMENTS_FILE, + 'abc.txt', 'def.txt']), + sorted(dependency.stage_job_resources( + options, + populate_requirements_cache=self.populate_requirements_cache))) + self.assertTrue( + os.path.isfile( + os.path.join(staging_dir, dependency.REQUIREMENTS_FILE))) + self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) + self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_requirements_file_not_present(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir @@ -159,16 +166,15 @@ class SetupTest(unittest.TestCase): '--requirements_file command line option.' % 'nosuchfile') def test_with_requirements_file_and_cache(self): - staging_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir self.update_options(options) options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, dependency.REQUIREMENTS_FILE) - options.view_as(SetupOptions).requirements_cache = os.path.join( - tempfile.gettempdir(), 'alternative-cache-dir') + options.view_as(SetupOptions).requirements_cache = self.make_temp_dir() self.create_temp_file( os.path.join(source_dir, dependency.REQUIREMENTS_FILE), 'nothing') self.assertEqual( @@ -184,8 +190,8 @@ class SetupTest(unittest.TestCase): self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_with_setup_file(self): - staging_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() self.create_temp_file( os.path.join(source_dir, 'setup.py'), 'notused') @@ -213,7 +219,7 @@ class SetupTest(unittest.TestCase): os.path.join(staging_dir, dependency.WORKFLOW_TARBALL_FILE))) def test_setup_file_not_present(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir @@ -228,8 +234,8 @@ class SetupTest(unittest.TestCase): '--setup_file command line option.' % 'nosuchfile') def test_setup_file_not_named_setup_dot_py(self): - staging_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(GoogleCloudOptions).staging_location = staging_dir @@ -279,7 +285,7 @@ class SetupTest(unittest.TestCase): return os.path.join(expected_to_folder, 'sdk-tarball') def test_sdk_location_default(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() expected_from_url = 'pypi' expected_from_path = self.override_pypi_download( expected_from_url, staging_dir) @@ -297,8 +303,8 @@ class SetupTest(unittest.TestCase): file_copy=dependency._dependency_file_copy)) def test_sdk_location_local(self): - staging_dir = tempfile.mkdtemp() - sdk_location = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + sdk_location = self.make_temp_dir() self.create_temp_file( os.path.join( sdk_location, @@ -319,7 +325,7 @@ class SetupTest(unittest.TestCase): self.assertEqual(f.read(), 'contents') def test_sdk_location_local_not_present(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() sdk_location = 'nosuchdir' with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() @@ -335,7 +341,7 @@ class SetupTest(unittest.TestCase): cm.exception.message) def test_sdk_location_gcs(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() sdk_location = 'gs://my-gcs-bucket/tarball.tar.gz' self.override_file_copy(sdk_location, staging_dir) @@ -349,8 +355,8 @@ class SetupTest(unittest.TestCase): dependency.stage_job_resources(options)) def test_with_extra_packages(self): - staging_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() self.create_temp_file( os.path.join(source_dir, 'abc.tar.gz'), 'nothing') self.create_temp_file( @@ -399,7 +405,7 @@ class SetupTest(unittest.TestCase): self.assertEqual(['gs://my-gcs-bucket/gcs.tar.gz'], gcs_copied_files) def test_with_extra_packages_missing_files(self): - staging_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() @@ -414,8 +420,8 @@ class SetupTest(unittest.TestCase): '--extra_packages command line option.' % 'nosuchfile.tar.gz') def test_with_extra_packages_invalid_file_name(self): - staging_dir = tempfile.mkdtemp() - source_dir = tempfile.mkdtemp() + staging_dir = self.make_temp_dir() + source_dir = self.make_temp_dir() self.create_temp_file( os.path.join(source_dir, 'abc.tgz'), 'nothing') with self.assertRaises(RuntimeError) as cm: