def testParseRequirementsFileInGcs(self):
"""Tests parsing requirements file in GCS file path."""
requirements_file = 'gs://mybucket/pypi.txt'
read_object_mock = self.StartObjectPatch(
storage_api.StorageClient, 'ReadObject')
command_util.ParseRequirementsFile(requirements_file)
config_object = storage_util.ObjectReference.FromUrl(requirements_file)
read_object_mock.assert_called_once_with(config_object)
def ConstructPatch(env_ref=None,
node_count=None,
update_pypi_packages_from_file=None,
clear_pypi_packages=None,
remove_pypi_packages=None,
update_pypi_packages=None,
clear_labels=None,
remove_labels=None,
update_labels=None,
clear_airflow_configs=None,
remove_airflow_configs=None,
update_airflow_configs=None,
clear_env_variables=None,
remove_env_variables=None,
update_env_variables=None,
update_image_version=None,
update_web_server_access_control=None,
cloud_sql_machine_type=None,
web_server_machine_type=None,
scheduler_cpu=None,
worker_cpu=None,
scheduler_memory_gb=None,
worker_memory_gb=None,
min_workers=None,
max_workers=None,
maintenance_window_start=None,
maintenance_window_end=None,
maintenance_window_recurrence=None,
release_track=base.ReleaseTrack.GA):
"""Constructs an environment patch.
Args:
env_ref: resource argument, Environment resource argument for environment
being updated.
node_count: int, the desired node count
update_pypi_packages_from_file: str, path to local requirements file
containing desired pypi dependencies.
clear_pypi_packages: bool, whether to uninstall all PyPI packages.
remove_pypi_packages: iterable(string), Iterable of PyPI packages to
uninstall.
update_pypi_packages: {string: string}, dict mapping PyPI package name to
extras and version specifier.
clear_labels: bool, whether to clear the labels dictionary.
remove_labels: iterable(string), Iterable of label names to remove.
update_labels: {string: string}, dict of label names and values to set.
clear_airflow_configs: bool, whether to clear the Airflow configs
dictionary.
remove_airflow_configs: iterable(string), Iterable of Airflow config
property names to remove.
update_airflow_configs: {string: string}, dict of Airflow config property
names and values to set.
clear_env_variables: bool, whether to clear the environment variables
dictionary.
remove_env_variables: iterable(string), Iterable of environment variables
to remove.
update_env_variables: {string: string}, dict of environment variable
names and values to set.
update_image_version: string, image version to use for environment upgrade
update_web_server_access_control: [{string: string}], Webserver access
control to set
cloud_sql_machine_type: str or None, Cloud SQL machine type used by the
Airflow database.
web_server_machine_type: str or None, machine type used by the Airflow web
server
scheduler_cpu: float or None, CPU allocated to Airflow scheduler.
Can be specified only in Composer 2.0.0.
worker_cpu: float or None, CPU allocated to each Airflow worker.
Can be specified only in Composer 2.0.0.
scheduler_memory_gb: float or None, memory allocated to Airflow scheduler.
Can be specified only in Composer 2.0.0.
worker_memory_gb: float or None, memory allocated to each Airflow worker.
Can be specified only in Composer 2.0.0.
min_workers: int or None, minimum number of workers
in the Environment. Can be specified only in Composer 2.0.0.
max_workers: int or None, maximumn number of workers
in the Environment. Can be specified only in Composer 2.0.0.
maintenance_window_start: Datetime or None, a starting date of the
maintenance window.
maintenance_window_end: Datetime or None, an ending date of the maintenance
window.
maintenance_window_recurrence: str or None, recurrence RRULE for the
maintenance window.
release_track: base.ReleaseTrack, the release track of command. Will dictate
which Composer client library will be used.
Returns:
(str, Environment), the field mask and environment to use for update.
Raises:
command_util.Error: if no update type is specified
"""
if node_count:
return _ConstructNodeCountPatch(node_count,
release_track=release_track)
if update_pypi_packages_from_file:
return _ConstructPyPiPackagesPatch(
True, [],
command_util.ParseRequirementsFile(update_pypi_packages_from_file),
release_track=release_track)
if clear_pypi_packages or remove_pypi_packages or update_pypi_packages:
return _ConstructPyPiPackagesPatch(clear_pypi_packages,
remove_pypi_packages,
update_pypi_packages,
release_track=release_track)
if clear_labels or remove_labels or update_labels:
return _ConstructLabelsPatch(clear_labels,
remove_labels,
update_labels,
release_track=release_track)
if (clear_airflow_configs or remove_airflow_configs
or update_airflow_configs):
return _ConstructAirflowConfigsPatch(clear_airflow_configs,
remove_airflow_configs,
update_airflow_configs,
release_track=release_track)
if clear_env_variables or remove_env_variables or update_env_variables:
return _ConstructEnvVariablesPatch(env_ref,
clear_env_variables,
remove_env_variables,
update_env_variables,
release_track=release_track)
if update_image_version:
return _ConstructImageVersionPatch(update_image_version,
release_track=release_track)
if update_web_server_access_control is not None:
return _ConstructWebServerAccessControlPatch(
update_web_server_access_control, release_track=release_track)
if cloud_sql_machine_type:
return _ConstructCloudSqlMachineTypePatch(cloud_sql_machine_type,
release_track=release_track)
if web_server_machine_type:
return _ConstructWebServerMachineTypePatch(web_server_machine_type,
release_track=release_track)
if (scheduler_cpu or worker_cpu or scheduler_memory_gb or worker_memory_gb
or min_workers or max_workers):
return _ConstructAutoscalingPatch(
scheduler_cpu=scheduler_cpu,
worker_cpu=worker_cpu,
scheduler_memory_gb=scheduler_memory_gb,
worker_memory_gb=worker_memory_gb,
worker_min_count=min_workers,
worker_max_count=max_workers,
release_track=release_track)
if maintenance_window_start and maintenance_window_end and maintenance_window_recurrence:
return _ConstructMaintenanceWindowPatch(maintenance_window_start,
maintenance_window_end,
maintenance_window_recurrence,
release_track=release_track)
raise command_util.Error(
'Cannot update Environment with no update type specified.')
def ConstructPatch(env_ref=None,
node_count=None,
update_pypi_packages_from_file=None,
clear_pypi_packages=None,
remove_pypi_packages=None,
update_pypi_packages=None,
clear_labels=None,
remove_labels=None,
update_labels=None,
clear_airflow_configs=None,
remove_airflow_configs=None,
update_airflow_configs=None,
clear_env_variables=None,
remove_env_variables=None,
update_env_variables=None,
update_image_version=None,
release_track=base.ReleaseTrack.GA):
"""Constructs an environment patch.
Args:
env_ref: resource argument, Environment resource argument for environment
being updated.
node_count: int, the desired node count
update_pypi_packages_from_file: str, path to local requirements file
containing desired pypi dependencies.
clear_pypi_packages: bool, whether to uninstall all PyPI packages.
remove_pypi_packages: iterable(string), Iterable of PyPI packages to
uninstall.
update_pypi_packages: {string: string}, dict mapping PyPI package name to
extras and version specifier.
clear_labels: bool, whether to clear the labels dictionary.
remove_labels: iterable(string), Iterable of label names to remove.
update_labels: {string: string}, dict of label names and values to set.
clear_airflow_configs: bool, whether to clear the Airflow configs
dictionary.
remove_airflow_configs: iterable(string), Iterable of Airflow config
property names to remove.
update_airflow_configs: {string: string}, dict of Airflow config property
names and values to set.
clear_env_variables: bool, whether to clear the environment variables
dictionary.
remove_env_variables: iterable(string), Iterable of environment variables
to remove.
update_env_variables: {string: string}, dict of environment variable
names and values to set.
update_image_version: string, image version to use for environment upgrade
release_track: base.ReleaseTrack, the release track of command. Will dictate
which Composer client library will be used.
Returns:
(str, Environment), the field mask and environment to use for update.
Raises:
command_util.Error: if no update type is specified
"""
if node_count:
return _ConstructNodeCountPatch(node_count, release_track=release_track)
if update_pypi_packages_from_file:
return _ConstructPyPiPackagesPatch(
True, [],
command_util.ParseRequirementsFile(update_pypi_packages_from_file),
release_track=release_track)
if clear_pypi_packages or remove_pypi_packages or update_pypi_packages:
return _ConstructPyPiPackagesPatch(
clear_pypi_packages,
remove_pypi_packages,
update_pypi_packages,
release_track=release_track)
if clear_labels or remove_labels or update_labels:
return _ConstructLabelsPatch(
clear_labels, remove_labels, update_labels, release_track=release_track)
if (clear_airflow_configs or remove_airflow_configs or
update_airflow_configs):
return _ConstructAirflowConfigsPatch(
clear_airflow_configs,
remove_airflow_configs,
update_airflow_configs,
release_track=release_track)
if clear_env_variables or remove_env_variables or update_env_variables:
return _ConstructEnvVariablesPatch(
env_ref,
clear_env_variables,
remove_env_variables,
update_env_variables,
release_track=release_track)
if update_image_version:
return _ConstructImageVersionPatch(
update_image_version, release_track=release_track)
raise command_util.Error(
'Cannot update Environment with no update type specified.')
def testParseRequirementsFileIgnoreWhitespaceLines(self):
"""Tests parsing requirements file skips lines with only whitespace."""
requirements_file = self.Touch(
self.root_path, contents='\n \npackage\n \t\npackage2[extra1]\n \n')
actual_entries = command_util.ParseRequirementsFile(requirements_file)
self.assertEqual({'package': '', 'package2': '[extra1]'}, actual_entries)
def testParseRequirementsFileStripsLines(self):
"""Tests parsing requirements file with leading and trailing whitespace."""
requirements_file = self.Touch(
self.root_path, contents=' package \n package2[extra1]\t\n')
actual_entries = command_util.ParseRequirementsFile(requirements_file)
self.assertEqual({'package': '', 'package2': '[extra1]'}, actual_entries)
def testParseRequirementsFileDuplicateKey(self):
"""Tests error when parsing pypi requirements with duplicate package."""
requirements_file = self.Touch(
self.root_path, contents='package>0.1\npackage\n')
with self.assertRaises(command_util.Error):
command_util.ParseRequirementsFile(requirements_file)
def testParseRequirementsFileNoVersionMatch(self):
"""Tests parsing pypi requirements without a version."""
requirements_file = self.Touch(self.root_path, contents='numpy')
actual_entries = command_util.ParseRequirementsFile(requirements_file)
self.assertEqual({'numpy': ''}, actual_entries)