<s> """<STR_LIT>""" <EOL> import logging <EOL> import os <EOL> import shutil <EOL> import tempfile <EOL> import unittest <EOL> from google . cloud . dataflow import utils <EOL> from google . cloud . dataflow . utils import dependency <EOL> from google . cloud . dataflow . utils import names <EOL> from google . cloud . dataflow . utils . options import GoogleCloudOptions <EOL> from google . cloud . dataflow . utils . options import PipelineOptions <EOL> from google . cloud . dataflow . utils . options import SetupOptions <EOL> from google . cloud . dataflow . version import __version__ <EOL> class SetupTest ( unittest . TestCase ) : <EOL> def update_options ( self , options ) : <EOL> setup_options = options . view_as ( SetupOptions ) <EOL> setup_options . sdk_location = '<STR_LIT>' <EOL> google_cloud_options = options . view_as ( GoogleCloudOptions ) <EOL> if google_cloud_options . temp_location is None : <EOL> google_cloud_options . temp_location = google_cloud_options . staging_location <EOL> def create_temp_file ( self , path , contents ) : <EOL> with open ( path , '<STR_LIT:w>' ) as f : <EOL> f . write ( contents ) <EOL> return f . name <EOL> def test_no_staging_location ( self ) : <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> dependency . stage_job_resources ( PipelineOptions ( ) ) <EOL> self . assertEqual ( '<STR_LIT>' , <EOL> cm . exception . message ) <EOL> def test_no_temp_location ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> google_cloud_options = options . view_as ( GoogleCloudOptions ) <EOL> google_cloud_options . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> google_cloud_options . temp_location = None <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( '<STR_LIT>' , <EOL> cm . exception . message ) <EOL> def test_no_main_session ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> options . view_as ( SetupOptions ) . save_main_session = False <EOL> self . update_options ( options ) <EOL> self . assertEqual ( <EOL> [ ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> def test_default_resources ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> self . assertEqual ( <EOL> [ names . PICKLED_MAIN_SESSION_FILE ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> self . assertTrue ( <EOL> os . path . isfile ( <EOL> os . path . join ( staging_dir , names . PICKLED_MAIN_SESSION_FILE ) ) ) <EOL> def test_with_requirements_file ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> source_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . requirements_file = os . path . join ( <EOL> source_dir , dependency . REQUIREMENTS_FILE ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , dependency . REQUIREMENTS_FILE ) , '<STR_LIT>' ) <EOL> self . assertEqual ( <EOL> [ dependency . REQUIREMENTS_FILE , <EOL> names . PICKLED_MAIN_SESSION_FILE ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> self . assertTrue ( <EOL> os . path . isfile ( <EOL> os . path . join ( staging_dir , dependency . REQUIREMENTS_FILE ) ) ) <EOL> def test_requirements_file_not_present ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . requirements_file = '<STR_LIT>' <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( <EOL> cm . exception . message , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % '<STR_LIT>' ) <EOL> def test_with_setup_file ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> source_dir = tempfile . mkdtemp ( ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . setup_file = os . path . join ( <EOL> source_dir , '<STR_LIT>' ) <EOL> self . assertEqual ( <EOL> [ dependency . WORKFLOW_TARBALL_FILE , <EOL> names . PICKLED_MAIN_SESSION_FILE ] , <EOL> dependency . stage_job_resources ( <EOL> options , <EOL> build_setup_args = [ <EOL> '<STR_LIT>' , '<STR_LIT:-c>' , '<STR_LIT>' , <EOL> os . path . join ( source_dir , dependency . WORKFLOW_TARBALL_FILE ) ] , <EOL> temp_dir = source_dir ) ) <EOL> self . assertTrue ( <EOL> os . path . isfile ( <EOL> os . path . join ( staging_dir , dependency . WORKFLOW_TARBALL_FILE ) ) ) <EOL> def test_setup_file_not_present ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . setup_file = '<STR_LIT>' <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( <EOL> cm . exception . message , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % '<STR_LIT>' ) <EOL> def test_setup_file_not_named_setup_dot_py ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> source_dir = tempfile . mkdtemp ( ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . setup_file = ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertTrue ( <EOL> cm . exception . message . startswith ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) ) <EOL> def override_file_copy ( self , expected_from_path , expected_to_dir ) : <EOL> def file_copy ( from_path , to_path ) : <EOL> if not from_path . endswith ( names . PICKLED_MAIN_SESSION_FILE ) : <EOL> self . assertEqual ( expected_from_path , from_path ) <EOL> self . assertEqual ( utils . path . join ( expected_to_dir , <EOL> names . DATAFLOW_SDK_TARBALL_FILE ) , <EOL> to_path ) <EOL> if from_path . startswith ( '<STR_LIT>' ) or to_path . startswith ( '<STR_LIT>' ) : <EOL> logging . info ( '<STR_LIT>' , from_path , to_path ) <EOL> else : <EOL> shutil . copyfile ( from_path , to_path ) <EOL> dependency . _dependency_file_copy = file_copy <EOL> def override_file_download ( self , expected_from_url , expected_to_folder ) : <EOL> def file_download ( from_url , _ ) : <EOL> self . assertEqual ( expected_from_url , from_url ) <EOL> tarball_path = os . path . join ( expected_to_folder , '<STR_LIT>' ) <EOL> with open ( tarball_path , '<STR_LIT:w>' ) as f : <EOL> f . write ( '<STR_LIT>' ) <EOL> return tarball_path <EOL> dependency . _dependency_file_download = file_download <EOL> return os . path . join ( expected_to_folder , '<STR_LIT>' ) <EOL> def test_sdk_location_default ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> expected_from_url = '<STR_LIT>' % ( <EOL> dependency . PACKAGES_URL_PREFIX , __version__ ) <EOL> expected_from_path = self . override_file_download ( <EOL> expected_from_url , staging_dir ) <EOL> self . override_file_copy ( expected_from_path , staging_dir ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . sdk_location = '<STR_LIT:default>' <EOL> self . assertEqual ( <EOL> [ names . PICKLED_MAIN_SESSION_FILE , <EOL> names . DATAFLOW_SDK_TARBALL_FILE ] , <EOL> dependency . stage_job_resources ( <EOL> options , <EOL> file_copy = dependency . _dependency_file_copy ) ) <EOL> def test_sdk_location_local ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> sdk_location = tempfile . mkdtemp ( ) <EOL> self . create_temp_file ( <EOL> os . path . join ( <EOL> sdk_location , <EOL> names . DATAFLOW_SDK_TARBALL_FILE ) , <EOL> '<STR_LIT>' ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . sdk_location = sdk_location <EOL> self . assertEqual ( <EOL> [ names . PICKLED_MAIN_SESSION_FILE , <EOL> names . DATAFLOW_SDK_TARBALL_FILE ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> tarball_path = os . path . join ( <EOL> staging_dir , names . DATAFLOW_SDK_TARBALL_FILE ) <EOL> with open ( tarball_path ) as f : <EOL> self . assertEqual ( f . read ( ) , '<STR_LIT>' ) <EOL> def test_sdk_location_local_not_present ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> sdk_location = '<STR_LIT>' <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . sdk_location = sdk_location <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % <EOL> sdk_location , <EOL> cm . exception . message ) <EOL> def test_sdk_location_gcs ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> sdk_location = '<STR_LIT>' <EOL> self . override_file_copy ( sdk_location , staging_dir ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . sdk_location = sdk_location <EOL> self . assertEqual ( <EOL> [ names . PICKLED_MAIN_SESSION_FILE , <EOL> names . DATAFLOW_SDK_TARBALL_FILE ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> def test_with_extra_packages ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> source_dir = tempfile . mkdtemp ( ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , dependency . EXTRA_PACKAGES_FILE ) , '<STR_LIT>' ) <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . extra_packages = [ <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , <EOL> '<STR_LIT>' ] <EOL> gcs_copied_files = [ ] <EOL> def file_copy ( from_path , to_path ) : <EOL> if from_path . startswith ( '<STR_LIT>' ) : <EOL> gcs_copied_files . append ( from_path ) <EOL> _ , from_name = os . path . split ( from_path ) <EOL> self . create_temp_file ( os . path . join ( to_path , from_name ) , '<STR_LIT>' ) <EOL> logging . info ( '<STR_LIT>' , from_path , to_path ) <EOL> elif to_path . startswith ( '<STR_LIT>' ) : <EOL> logging . info ( '<STR_LIT>' , from_path , to_path ) <EOL> else : <EOL> shutil . copyfile ( from_path , to_path ) <EOL> dependency . _dependency_file_copy = file_copy <EOL> self . assertEqual ( <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> dependency . EXTRA_PACKAGES_FILE , <EOL> names . PICKLED_MAIN_SESSION_FILE ] , <EOL> dependency . stage_job_resources ( options ) ) <EOL> with open ( os . path . join ( staging_dir , dependency . EXTRA_PACKAGES_FILE ) ) as f : <EOL> self . assertEqual ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> f . readlines ( ) ) <EOL> self . assertEqual ( [ '<STR_LIT>' ] , gcs_copied_files ) <EOL> def test_with_extra_packages_missing_files ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . extra_packages = [ '<STR_LIT>' ] <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( <EOL> cm . exception . message , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % '<STR_LIT>' ) <EOL> def test_with_extra_packages_invalid_file_name ( self ) : <EOL> staging_dir = tempfile . mkdtemp ( ) <EOL> source_dir = tempfile . mkdtemp ( ) <EOL> self . create_temp_file ( <EOL> os . path . join ( source_dir , '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> with self . assertRaises ( RuntimeError ) as cm : <EOL> options = PipelineOptions ( ) <EOL> options . view_as ( GoogleCloudOptions ) . staging_location = staging_dir <EOL> self . update_options ( options ) <EOL> options . view_as ( SetupOptions ) . extra_packages = [ <EOL> os . path . join ( source_dir , '<STR_LIT>' ) ] <EOL> dependency . stage_job_resources ( options ) <EOL> self . assertEqual ( <EOL> cm . exception . message , <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' % os . path . join ( source_dir , '<STR_LIT>' ) ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> logging . getLogger ( ) . setLevel ( logging . INFO ) <EOL> unittest . main ( ) </s>
