diff --git a/geos_ats_package/geos_ats/baseline_io.py b/geos_ats_package/geos_ats/baseline_io.py index d554800..c8e7127 100644 --- a/geos_ats_package/geos_ats/baseline_io.py +++ b/geos_ats_package/geos_ats/baseline_io.py @@ -2,6 +2,8 @@ import logging import tempfile import shutil +import yaml +import time from google.cloud import storage logger = logging.getLogger( 'geos_ats' ) @@ -42,16 +44,14 @@ def download_baselines( bucket_name, blob_name, baseline_path, force_redownload= logger.info( 'Downloading baselines...' ) tmpdir = tempfile.TemporaryDirectory() archive_name = os.path.join( tmpdir.name, 'baselines.tar.gz' ) - bucket = storage_client.bucket( bucket_name ) + client = storage.Client() + bucket = client.bucket( bucket_name ) blob = bucket.blob( blob_name ) blob.download_to_filename( archive_name ) # Unpack new baselines logger.info( 'Unpacking baselines...' ) shutil.unpack_archive( archive_name, baseline_path, format='gztar' ) - with open( status_path, 'w' ) as f: - f.write( blob_name ) - logger.info( 'Finished fetching baselines!' ) except Exception as e: @@ -92,7 +92,8 @@ def upload_baselines( bucket_name, blob_name, baseline_path ): shutil.make_archive( archive_name, format='gztar', base_dir=baseline_path ) logger.info( 'Uploading baseline files...' ) - bucket = storage_client.bucket( bucket_name ) + client = storage.Client() + bucket = client.bucket( bucket_name ) blob = bucket.blob( blob_name ) blob.upload_from_filename( archive_name, if_generation_match=0 ) logger.info( 'Finished uploading baselines!' ) @@ -104,3 +105,49 @@ def upload_baselines( bucket_name, blob_name, baseline_path ): # Reset the local blob name with open( status_path, 'w' ) as f: f.write( last_blob_name ) + + +def manage_baselines( options ): + """ + Manage the integrated test baselines + """ + # Check for integrated test yaml file + test_yaml = '' + if options.integrateTestsYAML: + test_yaml = options.integrateTestsYAML + else: + os.path.join( options.geos_bin_dir ), '..', '..', '.integrated_tests.yaml' + + if not os.path.isfile( test_yaml ): + raise Exception( f'Could not find the integrated test yaml file: {test_yaml}' ) + + test_options = yaml.load( open( test_yaml ) ) + baseline_options = test_options.get( 'baselines', {} ) + for k in [ 'bucket', 'latest' ]: + if k not in baseline_options: + raise Exception( + f'Required information (baselines/{k}) missing from integrated test yaml file: {test_yaml}' ) + + # Manage baselines + if options.action == 'upload_baselines': + if os.path.isdir( options.baselineDir ): + epoch = int( time.time() ) + upload_name = f'integrated_test_baseline_{epoch}.tar.gz' + upload_baselines( baseline_options[ 'bucket' ], upload_name, options.baselineDir ) + + # Update the test config file + baseline_options[ 'latest' ] = upload_name + with open( test_yaml, 'w' ) as f: + yaml.dump( baseline_options, f ) + quit() + else: + raise Exception( f'Could not find the requested baselines to upload: {options.baselineDir}' ) + + download_baselines( baseline_options[ 'bucket' ], baseline_options[ 'latest' ], options.baselineDir ) + + # Cleanup + if not os.path.isdir( options.baselineDir ): + raise Exception( f'Could not find the specified baseline directory: {options.baselineDir}' ) + + if options.action == 'download_baselines': + quit() diff --git a/geos_ats_package/geos_ats/command_line_parsers.py b/geos_ats_package/geos_ats/command_line_parsers.py index 09ead6f..d39720f 100644 --- a/geos_ats_package/geos_ats/command_line_parsers.py +++ b/geos_ats_package/geos_ats/command_line_parsers.py @@ -17,6 +17,8 @@ "rebaseline": "rebaseline the testcases from a previous run.", "rebaselinefailed": "rebaseline only failed testcases from a previous run.", "report": "generate a text or html report, see config for the reporting options.", + "upload_baselines": "Upload baselines to bucket", + "download_baselines": "Download baselines from bucket", } check_options = { @@ -46,6 +48,8 @@ def build_command_line_parser(): parser.add_argument( "-b", "--baselineDir", type=str, help="Root baseline directory" ) + parser.add_argument( "-y", "--integrateTestsYAML", type=str, help="Baseline bucket name", default='' ) + action_names = ','.join( action_options.keys() ) parser.add_argument( "-a", "--action", type=str, default="run", help=f"Test actions options ({action_names})" ) diff --git a/geos_ats_package/geos_ats/main.py b/geos_ats_package/geos_ats/main.py index 42e43c1..5de47bc 100644 --- a/geos_ats_package/geos_ats/main.py +++ b/geos_ats_package/geos_ats/main.py @@ -5,7 +5,7 @@ import subprocess import time import logging -from geos_ats import command_line_parsers +from geos_ats import command_line_parsers, baseline_io test_actions = ( "run", "rerun", "check", "continue" ) report_actions = ( "run", "rerun", "report", "continue" ) @@ -292,6 +292,7 @@ def main(): os.chdir( ats_root_dir ) os.makedirs( options.workingDir, exist_ok=True ) create_log_directory( options ) + baseline_io.manage_baselines( options ) # Check the test configuration from geos_ats import configuration_record diff --git a/geos_ats_package/setup.cfg b/geos_ats_package/setup.cfg index 2d64f43..d2e7859 100644 --- a/geos_ats_package/setup.cfg +++ b/geos_ats_package/setup.cfg @@ -17,6 +17,8 @@ install_requires = numpy lxml tabulate + pyyaml + google-cloud ats @ https://github.com/LLNL/ATS/archive/refs/tags/7.0.105.tar.gz python_requires = >=3.7