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)