def tearDown(self):
self.delete_gcs_bucket(BUCKET)
with provide_gcp_context(GCP_BIGQUERY_KEY, scopes=SCOPES):
hook = BigQueryHook()
hook.delete_dataset(dataset_id='airflow_test',
delete_contents=True)
super().tearDown()
def assert_remote_logs(self, expected_message, ti):
with provide_gcp_context(GCP_STACKDDRIVER), conf_vars({
('logging', 'remote_logging'): 'True',
('logging', 'remote_base_log_folder'): f"stackdriver://{self.log_name}",
}):
from airflow.config_templates import airflow_local_settings
importlib.reload(airflow_local_settings)
settings.configure_logging()
task_log_reader = TaskLogReader()
logs = "\n".join(task_log_reader.read_log_stream(ti, try_number=None, metadata={}))
self.assertIn(expected_message, logs)
def assert_remote_logs(self, expected_message, ti):
with provide_gcp_context(GCP_GCS_KEY), conf_vars({
('logging', 'remote_logging'): 'True',
('logging', 'remote_base_log_folder'): f"gs://{self.bucket_name}/path/to/logs",
('logging', 'remote_log_conn_id'): "google_cloud_default",
}):
from airflow.config_templates import airflow_local_settings
importlib.reload(airflow_local_settings)
settings.configure_logging()
task_log_reader = TaskLogReader()
logs = "\n".join(task_log_reader.read_log_stream(ti, try_number=None, metadata={}))
self.assertIn(expected_message, logs)
def helper():
with provide_gcp_context(GCP_GCS_TRANSFER_KEY):
# Create buckets
GoogleSystemTest.create_gcs_bucket(GCP_TRANSFER_SECOND_TARGET_BUCKET, location="asia-east1")
GoogleSystemTest.create_gcs_bucket(GCP_TRANSFER_FIRST_TARGET_BUCKET)
GoogleSystemTest.upload_content_to_gcs("test_contents", GCP_TRANSFER_FIRST_TARGET_BUCKET, "test.txt")
# Grant bucket permissions
project_number = GoogleSystemTest.get_project_number(GCP_PROJECT_ID)
account_email = f"project-{project_number}@storage-transfer-service.iam.gserviceaccount.com"
GoogleSystemTest.grant_bucket_access(GCP_TRANSFER_FIRST_TARGET_BUCKET, account_email)
GoogleSystemTest.grant_bucket_access(GCP_TRANSFER_SECOND_TARGET_BUCKET, account_email)
yield
# Remove buckets
GoogleSystemTest.delete_gcs_bucket(GCP_TRANSFER_SECOND_TARGET_BUCKET)
GoogleSystemTest.delete_gcs_bucket(GCP_TRANSFER_FIRST_TARGET_BUCKET)