def test_execute(self, mock_make_query, mock_write_object_to_file, mock_upload):
mock_make_query.return_value = SALESFORCE_RESPONSE
operator = SalesforceToGcsOperator(
query=QUERY,
bucket_name=GCS_BUCKET,
object_name=GCS_OBJECT_PATH,
salesforce_conn_id=SALESFORCE_CONNECTION_ID,
gcp_conn_id=GCP_CONNECTION_ID,
include_deleted=INCLUDE_DELETED,
query_params=QUERY_PARAMS,
export_format="json",
coerce_to_timestamp=True,
record_time_added=True,
task_id=TASK_ID,
)
result = operator.execute({})
mock_make_query.assert_called_once_with(
query=QUERY, include_deleted=INCLUDE_DELETED, query_params=QUERY_PARAMS
)
mock_write_object_to_file.assert_called_once_with(
query_results=SALESFORCE_RESPONSE['records'],
filename=mock.ANY,
fmt="json",
coerce_to_timestamp=True,
record_time_added=True,
)
mock_upload.assert_called_once_with(
bucket_name=GCS_BUCKET, object_name=GCS_OBJECT_PATH, filename=mock.ANY, gzip=False
)
self.assertEqual(EXPECTED_GCS_URI, result)
) as dag:
create_bucket = GCSCreateBucketOperator(
task_id="create_bucket",
bucket_name=GCS_BUCKET,
project_id=GCP_PROJECT_ID,
gcp_conn_id=GCS_CONN_ID,
)
# [START howto_operator_salesforce_to_gcs]
gcs_upload_task = SalesforceToGcsOperator(
query=QUERY,
include_deleted=True,
bucket_name=GCS_BUCKET,
object_name=GCS_OBJ_PATH,
salesforce_conn_id=SALESFORCE_CONN_ID,
export_format='csv',
coerce_to_timestamp=False,
record_time_added=False,
gcp_conn_id=GCS_CONN_ID,
task_id="upload_to_gcs",
dag=dag,
)
# [END howto_operator_salesforce_to_gcs]
create_dataset = BigQueryCreateEmptyDatasetOperator(
task_id="create_dataset", dataset_id=DATASET_NAME, project_id=GCP_PROJECT_ID, gcp_conn_id=GCS_CONN_ID
)
create_table = BigQueryCreateEmptyTableOperator(
task_id="create_table",
dataset_id=DATASET_NAME,