def test_create_views(bq_client, cli_runner, test_data): name_mock = Mock() name_mock.name = 'id' bq_client().get_table.side_effect = [ Mock(table_type='TABLE', table_id='table1', schema=[name_mock, Mock()]), NotFound('table1 was not found'), Mock(table_type='TABLE', table_id='table2', schema=[name_mock, Mock()]), Mock(table_type='TABLE', table_id='table2', schema=[name_mock, Mock()]), Mock(table_type='TABLE', table_id='table3', schema=[name_mock, Mock()]), NotFound('table3 was not found'), Mock(table_type='TABLE', table_id='table4', schema=[name_mock, Mock()]), NotFound('table4 was not found'), Mock(table_type='TABLE', table_id='table5', schema=[name_mock, Mock()]), NotFound('table5 was not found'), Mock(table_type='TABLE', table_id='table6', schema=[name_mock, Mock()]), NotFound('table6 was not found'), ] result = cli_runner.invoke_and_assert_exit( 0, bq_admin, ['create-views', str(test_data.path('bigquery-view-specs.json'))]) assert result.stdout == """\
def get_by_attribute(self, attribute: str, value: Any) -> T: docs = self.query_by_attribute(attribute=attribute, value=value) if len(docs) == 0: raise NotFound( f"Document could not be found in {self.collection_name} with `{attribute}=={value}`" ) elif len(docs) > 1: raise NotFound( f"Founder more than one document in {self.collection_name} with `{attribute}=={value}`" ) return docs[0]
def test_limited_project_access(self, MockResourceManagerClient, MockLoggingClient): MockResourceManagerClient.return_value.list_projects.return_value = [ MagicMock(project_id='proj1'), MagicMock(project_id='proj2'), MagicMock(project_id='proj3'), ] MockLoggingClient.return_value.project = 'proj1' MockLoggingClient.return_value.list_entries.side_effect = [ PermissionDenied(''), iter(SAMPLE_ENTRIES), NotFound(''), ] reader = Reader( start_time=datetime(2018, 4, 3, 9, 51, 22), end_time=datetime(2018, 4, 3, 10, 51, 33), collect_multiple_projects=True, ) self.assertEqual( reader.log_list, [ BASE_LOG_NAME.format('proj1'), BASE_LOG_NAME.format('proj2'), BASE_LOG_NAME.format('proj3'), ], ) self.assertEqual(list(reader), [FlowRecord(x) for x in SAMPLE_ENTRIES])
def run(self): if self.opts.command == "create_key": if self.set_hash_key(): print(f"Saved secret at {self.secret_path}") else: print(f"Hash key already exists at {self.secret_path}") exit(0) os.environ["GCLOUD_PROJECT"] = self.opts.project os.system( f"gcloud alpha firestore databases create --project {self.opts.project} --region {self.opts.region}" ) db = firestore.Client() col = db.collection(self.opts.collection) key = self.get_hash_key() if key == None: raise NotFound( "Hash key does not exist yet. Please run `hasher.py create_key --secret $SECRET` first" ) if self.opts.command == "upload": print("Hashing and uploading SSNs...") for ssn in open(self.opts.infile): digest = self.hash_ssn(ssn, key) col.document(digest).set({u'exists': True}) print("Done!") elif self.opts.command == 'verify': print("Verifying and counting SSNs...") count = 0 for ssn in open(self.opts.infile): digest = self.hash_ssn(ssn, key) doc = col.document(digest).get() if doc.exists: count += 1 print(f"Found {count} valid SSNs")
def get_bucket(self, bucket_name): if bucket_name not in self.mock_gcs_fs: raise NotFound( 'GET https://www.googleapis.com/storage/v1/b/%s' '?projection=noAcl: Not Found' % bucket_name) return self.bucket(bucket_name)
def get( self, id: str, include_attributes: Optional[List[str]] = None, return_attribute: Optional[str] = None, ) -> T: if include_attributes is not None and return_attribute is not None: raise KeyError( f"Cannot provide both `include_attributes` and `return_attribute`" ) if return_attribute is not None: include_attributes = [return_attribute] if include_attributes is not None: if any((not self.has_attribute(attribute) for attribute in include_attributes)): raise KeyError("Invalid attribute provided: `{attribute}`") doc_ref = self.collection.document(id) doc = doc_ref.get(field_paths=include_attributes) if not doc.exists: raise NotFound( f"Document {self.collection_name}.{id} could not be found") if return_attribute is not None: return doc.get(return_attribute) return self.schema(**{**doc.to_dict(), "id": doc.id})
def test_execute_if_cluster_exists_in_deleting_state( self, mock_hook, mock_get_cluster, mock_create_cluster, mock_generator ): cluster = mock.MagicMock() cluster.status.state = 0 cluster.status.DELETING = 0 cluster2 = mock.MagicMock() cluster2.status.state = 0 cluster2.status.ERROR = 0 mock_create_cluster.side_effect = [AlreadyExists("test"), cluster2] mock_generator.return_value = [0] mock_get_cluster.side_effect = [cluster, NotFound("test")] op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_config=CONFIG, labels=LABELS, cluster_name=CLUSTER_NAME, delete_on_error=True, gcp_conn_id=GCP_CONN_ID, ) with self.assertRaises(AirflowException): op.execute(context={}) calls = [mock.call(mock_hook.return_value), mock.call(mock_hook.return_value)] mock_get_cluster.assert_has_calls(calls) mock_create_cluster.assert_has_calls(calls) mock_hook.return_value.diagnose_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME )
def mock_bigquery_client(monkeypatch): from google.api_core.exceptions import NotFound import google.cloud.bigquery import google.cloud.bigquery.table mock_client = mock.create_autospec(google.cloud.bigquery.Client) # Constructor returns the mock itself, so this mock can be treated as the # constructor or the instance. mock_client.return_value = mock_client mock_schema = [google.cloud.bigquery.SchemaField("_f0", "INTEGER")] # Mock out SELECT 1 query results. mock_query = mock.create_autospec(google.cloud.bigquery.QueryJob) mock_query.job_id = "some-random-id" mock_query.state = "DONE" mock_rows = mock.create_autospec(google.cloud.bigquery.table.RowIterator) mock_rows.total_rows = 1 mock_rows.schema = mock_schema mock_rows.__iter__.return_value = [(1, )] mock_query.result.return_value = mock_rows mock_client.query.return_value = mock_query # Mock table creation. mock_client.get_table.side_effect = NotFound("nope") monkeypatch.setattr(google.cloud.bigquery, "Client", mock_client) mock_client.reset_mock() return mock_client
def upload_file(file, folder, upload_bucket, make_public=True, is_dev=False): """ Uploads a file to a given Cloud Storage bucket and returns the public url to the new object. """ if not file: return None filename = _safe_filename(file.filename) if is_dev: file.save(os.path.join(upload_bucket, filename)) url = "/" + upload_bucket + "/" + filename else: client = storage.Client() bucket = client.bucket(upload_bucket) blob = bucket.blob(folder + filename) try: blob.upload_from_string(file.read(), content_type=file.content_type) except NotFound: raise NotFound( "Could not upload file, bucket or folder %s does not exist" % (upload_bucket + folder)) if make_public: blob.make_public() url = blob.public_url if type(url) is bytes: url = url.decode('utf-8') return url
def create_column_family(self, column_family_name, table_name, max_age=None, nr_max_versions=None, gc_rule_union=None): """Create a column family and add it to a table. Garbage collection rules can be included to the column family. Args: column_family_name (str): table_name (str): max_age (int): the time to live in days nr_max_versions (int): the number of versions that should be kept gc_rule_union (bool or None): if both max_age and nr_max_versions are specified, then this parameter should be a bool. If True, then the max age and the max versions rules are unified, if False, then the intersection of the rules is used. Returns: google.cloud.bigtable.column_family.ColumnFamily """ if max_age and nr_max_versions: # Both rules are specified, this also means a merge method must be specified (union or intersection) time_to_live = dt.timedelta(days=max_age) max_age_rule = bt_column_family.MaxAgeGCRule(time_to_live) max_versions_rule = bt_column_family.MaxVersionsGCRule( nr_max_versions) if gc_rule_union is None: raise Conflict( "If max_age and nr_max_versions are both specified, then gc_rule_union cannot be None." ) elif gc_rule_union: gc_rule = bt_column_family.GCRuleUnion( rules=[max_age_rule, max_versions_rule]) else: gc_rule = bt_column_family.GCRuleIntersection( rules=[max_age_rule, max_versions_rule]) elif max_age: # only max age is specified time_to_live = dt.timedelta(days=max_age) gc_rule = bt_column_family.MaxAgeGCRule(time_to_live) elif nr_max_versions: # only max number of versions is specified gc_rule = bt_column_family.MaxVersionsGCRule(nr_max_versions) else: # no rule is specified gc_rule = None table = self.instance.table(table_name) if not table.exists(): raise NotFound( "Table name '{}' does not exist.".format(table_name)) logging.info("Creating column family '%s' in table '%s'.", column_family_name, table_name) column_family = bt_column_family.ColumnFamily(column_family_name, table, gc_rule) column_family.create()
def test_get_missing_key(self, mock_get_credentials, mock_client): mock_client.secret_version_path.return_value = "full-path" mock_client.access_secret_version.side_effect = NotFound('test-msg') secrets_manager_hook = SecretsManagerHook(gcp_conn_id='test') mock_get_credentials.assert_called_once_with() secret = secrets_manager_hook.get_secret(secret_id="secret") mock_client.secret_version_path.assert_called_once_with('example-project', 'secret', 'latest') mock_client.access_secret_version.assert_called_once_with("full-path") self.assertIsNone(secret)
def download_as_string(self): fs = self.bucket.client.mock_gcs_fs try: return fs[self.bucket.name]['blobs'][self.name]['data'] except KeyError: raise NotFound('GET https://www.googleapis.com/download/storage' '/v1/b/%s/o/%s?alt=media: Not Found' % (self.bucket.name, self.name))
def test_get_variable_non_existent_key(self, mock_client_callable, mock_get_creds): mock_get_creds.return_value = CREDENTIALS, PROJECT_ID mock_client = mock.MagicMock() mock_client_callable.return_value = mock_client # The requested secret id or secret version does not exist mock_client.access_secret_version.side_effect = NotFound('test-msg') secrets_manager_backend = CloudSecretsManagerBackend(variables_prefix=VARIABLES_PREFIX) self.assertIsNone(secrets_manager_backend.get_variable(VAR_KEY))
def test_delete_cluster_not_found(self, wait_mock, convert_mock, log_mock): from google.api_core.exceptions import NotFound # To force an error message = 'Not Found' self.gke_hook._client.delete_cluster.side_effect = NotFound(message=message) self.gke_hook.delete_cluster('not-existing') wait_mock.assert_not_called() convert_mock.assert_not_called() log_mock.info.assert_any_call("Assuming Success: %s", message)
def delete(self): fs = self.bucket.client.mock_gcs_fs if (self.bucket.name not in fs or self.name not in fs[self.bucket.name]['blobs']): raise NotFound('DELETE https://www.googleapis.com/storage/v1/b' '/%s/o/%s: Not Found' % (self.bucket.name, self.name)) del fs[self.bucket.name]['blobs'][self.name]
def test_get_or_create_bucket_no_exist(mocker): mock_storage_client = MockStorageClient() mocker.patch.object(mock_storage_client, "get_bucket", side_effect=NotFound("mock message")) mocker.patch.object(mock_storage_client, "create_bucket") mock_storage_client.create_bucket.return_value = MockBucket(name="bar") bucket = cloud_build.get_or_create_bucket("bar", mock_storage_client) assert bucket.__dict__ == MockBucket(name="bar").__dict__
def test_upset_user_profile_in_firestore_not_existant(when, db): not_existant_up = Box({"sub": "non_existent"}) collection = mock(CollectionReference) document = mock(DocumentReference) when(db).collection("pax").thenReturn(collection) when(collection).document("non_existent").thenReturn(document) when(document).update(...).thenRaise(NotFound("not found")) when(document).set(...).thenReturn() import main main.upset_user_profile_in_firestore("non_existent", not_existant_up)
def delete_cluster(self, project_id, region, cluster_name): self._check_region_matches_endpoint(region) cluster_key = (project_id, region, cluster_name) cluster = self.mock_clusters.get(cluster_key) if not cluster: raise NotFound('Not found: Cluster ' + _cluster_path(*cluster_key)) cluster.status.state = _cluster_state_value('DELETING')
def test_get_conn_uri_non_existent_key(self, mock_client_callable, mock_get_creds): mock_get_creds.return_value = CREDENTIALS, PROJECT_ID mock_client = mock.MagicMock() mock_client_callable.return_value = mock_client # The requested secret id or secret version does not exist mock_client.access_secret_version.side_effect = NotFound('test-msg') secrets_manager_backend = CloudSecretsManagerBackend(connections_prefix=CONNECTIONS_PREFIX) self.assertIsNone(secrets_manager_backend.get_conn_uri(conn_id=CONN_ID)) self.assertEqual([], secrets_manager_backend.get_connections(conn_id=CONN_ID))
def upload_from_string(self, data): if self.bucket.name not in self._fs: raise NotFound('POST https://www.googleapis.com/upload/storage' '/v1/b/%s/o?uploadType=multipart: Not Found' % self.bucket.name) fs_objs = self._fs[self.bucket.name]['blobs'] fs_obj = fs_objs.setdefault(self.name, dict(data=b'')) fs_obj['data'] = data self._set_md5_hash()
def test_get_non_existing_key(self, mock_client_info, mock_secrets_client): mock_client = mock.MagicMock() mock_client_info.return_value = mock.MagicMock() mock_secrets_client.return_value = mock_client mock_client.secret_version_path.return_value = "full-path" # The requested secret id or secret version does not exist mock_client.access_secret_version.side_effect = NotFound('test-msg') secrets_client = _SecretManagerClient(credentials="credentials") secret = secrets_client.get_secret(secret_id="missing", project_id="project_id") mock_client.secret_version_path.assert_called_once_with("project_id", 'missing', 'latest') assert secret is None mock_client.access_secret_version.assert_called_once_with('full-path')
def list_blobs(self, prefix=None): fs = self.client.mock_gcs_fs if self.name not in fs: raise NotFound('GET https://www.googleapis.com/storage/v1/b' '/%s/o?projection=noAcl: Not Found' % self.name) for blob_name in sorted(fs[self.name]['blobs']): if prefix and not blob_name.startswith(prefix): continue yield self.blob(blob_name)
def get_cluster(self, project_id, region, cluster_name): self._check_region_matches_endpoint(region) cluster_key = (project_id, region, cluster_name) if cluster_key not in self.mock_clusters: raise NotFound('Not Found: Cluster ' + _cluster_path(*cluster_key)) cluster = self.mock_clusters[cluster_key] result = deepcopy(cluster) self._simulate_progress(project_id, region, cluster_name) return result
def download_as_string(self, client=None, start=None, end=None): try: data = self._fs[self.bucket.name]['blobs'][self.name]['data'] except KeyError: raise NotFound('GET %s?alt=media: Not Found' % self._blob_uri()) if start is not None and start >= len(data): # it doesn't care if *end* exceeds the range raise RequestRangeNotSatisfiable( 'GET %s?alt=media: Request range not satisfiable' % self._blob_uri()) return data[start:end]
def get_job(self, project_id, region, job_id): self._check_region_matches_endpoint(region) job_key = (project_id, region, job_id) job = self.mock_jobs.get(job_key) if not job: raise NotFound('Not found: Job ' + _job_path(*job_key)) result = deepcopy(job) self._simulate_progress(job) return result
def test_not_found_race_condition(self): # in some cases the blob for the log file appears but # raises NotFound when read from self.start( patch( 'tests.mock_google.storage.MockGoogleStorageBlob' '.download_as_string', side_effect=NotFound('race condition'))) log_uri = self.URI + '.000000000' self.append_data(log_uri, b'log line\nanother log line\n') self.assertEqual(self.get_new_lines(), [])
def _safe_lookup(self, project, dataset_id=None, table_id=None): # type: (str, Optional[str], Optional[str]) -> Any """Look up data in self._datasets, raise NotFound if the key(s) is/are not present. Can be used to look up a project, dataset, or table. Args: project: project to look up dataset_id: If not specified, return datasets in project. If specified, the dataset to look up. table_id: If not specified, return tables in datasets. If specified, the table to look up. Returns: The map of all datasets in a requested project OR the map of all tables in a requested dataset OR a specific table depending on the specificity of the arguments. """ project_map = self._datasets # First look up the project if project not in project_map: raise NotFound("Project {} not found".format(project)) dataset_map = project_map[project] # If requested, look up the dataset if dataset_id is None: return dataset_map if dataset_id not in dataset_map: raise NotFound("Dataset {} not found".format(dataset_id)) table_map = dataset_map[dataset_id] # If requested, look up a table. if table_id is None: return table_map if table_id not in table_map: raise NotFound("Table {} not found".format(table_id)) return table_map[table_id]
def test_delete_photos_from_bucket_when_bucket_not_found(): mock_storage_client = mock.MagicMock() mock_bucket = mock.MagicMock() mock_storage_client.get_bucket.return_value = mock_bucket mock_bucket.delete_blob.side_effect = NotFound("message") with mock.patch("sys.stdout", new_callable=io.StringIO) as mock_stdout: utils.delete_photos_from_bucket( mock_storage_client, "manuvic/photos/transporteur-de-personne-trapu020", "not-existing-bucket") assert "manuvic/photos/transporteur-de-personne-trapu020 not found" in mock_stdout.getvalue( )
def test_get_variable_non_existent_key(self, mock_client_callable, mock_get_creds): mock_get_creds.return_value = CREDENTIALS, PROJECT_ID mock_client = mock.MagicMock() mock_client_callable.return_value = mock_client # The requested secret id or secret version does not exist mock_client.access_secret_version.side_effect = NotFound('test-msg') secrets_manager_backend = CloudSecretManagerBackend(variables_prefix=VARIABLES_PREFIX) secret_id = secrets_manager_backend.build_path(VARIABLES_PREFIX, VAR_KEY, SEP) with self.assertLogs(secrets_manager_backend.client.log, level="ERROR") as log_output: self.assertIsNone(secrets_manager_backend.get_variable(VAR_KEY)) self.assertRegex( log_output.output[0], f"Google Cloud API Call Error \\(NotFound\\): Secret ID {secret_id} not found", )
def find_bucket_with_prefix(bucket_iter: Iterator[Bucket], prefix: str) -> Union[str, None]: """"Finds bucket in a project based on bucket prefix. Args: bucket_iter: Iterator of google.cloud.storage.Bucket instances prefix: Bucket name prefix to search for Returns: Bucket name with the specified prefix. """ for bucket in bucket_iter: if bucket.name.startswith(prefix): return bucket.name raise NotFound(f'No bucket found with prefix: {prefix}')