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}')
