def test_new_attributes_cannot_be_set_on_child_of_seal(self):
m = mock.Mock()
m.test.test2 = 1
mock.seal(m)
with self.assertRaises(AttributeError):
m.test.test3 = 1
def test_existing_attributes_allowed_after_seal(self):
m = mock.Mock()
m.test.return_value = 3
mock.seal(m)
self.assertEqual(m.test(), 3)
def test_existing_attributes_can_be_set_on_seal(self):
m = mock.Mock()
m.test.test2 = 1
mock.seal(m)
m.test.test2 = 2
self.assertEqual(m.test.test2, 2)
def test_sealed_exception_has_attribute_name(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.SECRETE_name
self.assertIn("SECRETE_name", str(cm.exception))
def test_existing_attributes_can_be_set_on_seal(self):
m = mock.Mock()
m.test.test2 = 1
mock.seal(m)
m.test.test2 = 2
self.assertEqual(m.test.test2, 2)
def test_sealed_exception_has_attribute_name(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.SECRETE_name
self.assertIn("SECRETE_name", str(cm.exception))
def test_existing_attributes_allowed_after_seal(self):
m = mock.Mock()
m.test.return_value = 3
mock.seal(m)
self.assertEqual(m.test(), 3)
def test_new_attributes_cannot_be_set_on_child_of_seal(self):
m = mock.Mock()
m.test.test2 = 1
mock.seal(m)
with self.assertRaises(AttributeError):
m.test.test3 = 1
def test_call_chain_is_maintained(self):
m = mock.Mock()
m.test1().test2.test3().test4
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.test1().test2.test3().test4()
self.assertIn("mock.test1().test2.test3().test4", str(cm.exception))
def test_new_attributes_cannot_be_accessed_on_seal(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError):
m.test
with self.assertRaises(AttributeError):
m()
def test_call_chain_is_maintained(self):
m = mock.Mock()
m.test1().test2.test3().test4
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.test1().test2.test3().test4()
self.assertIn("mock.test1().test2.test3().test4", str(cm.exception))
def test_attribute_chain_is_maintained(self):
m = mock.Mock(name="mock_name")
m.test1.test2.test3.test4
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.test1.test2.test3.test4.boom
self.assertIn("mock_name.test1.test2.test3.test4.boom", str(cm.exception))
def test_new_attributes_cannot_be_accessed_on_seal(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError):
m.test
with self.assertRaises(AttributeError):
m()
def test_attribute_chain_is_maintained(self):
m = mock.Mock(name="mock_name")
m.test1.test2.test3.test4
mock.seal(m)
with self.assertRaises(AttributeError) as cm:
m.test1.test2.test3.test4.boom
self.assertIn("mock_name.test1.test2.test3.test4.boom",
str(cm.exception))
def test_seals_dont_recurse_on_manual_attributes(self):
m = mock.Mock(name="root_mock")
m.test1.test2 = mock.Mock(name="not_sealed")
m.test1.test2.test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2.test3, 4)
m.test1.test2.test4 # Does not raise
m.test1.test2.test4 = 1 # Does not raise
def test_integration_with_spec_method_definition(self):
"""You need to defin the methods, even if they are in the spec"""
m = mock.Mock(SampleObject)
m.method_sample1.return_value = 1
mock.seal(m)
self.assertEqual(m.method_sample1(), 1)
with self.assertRaises(AttributeError):
m.method_sample2()
def test_seals_dont_recurse_on_manual_attributes(self):
m = mock.Mock(name="root_mock")
m.test1.test2 = mock.Mock(name="not_sealed")
m.test1.test2.test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2.test3, 4)
m.test1.test2.test4 # Does not raise
m.test1.test2.test4 = 1 # Does not raise
def test_integration_with_spec_method_definition(self):
"""You need to defin the methods, even if they are in the spec"""
m = mock.Mock(SampleObject)
m.method_sample1.return_value = 1
mock.seal(m)
self.assertEqual(m.method_sample1(), 1)
with self.assertRaises(AttributeError):
m.method_sample2()
def test_extract_error(self, mock_client, mock_hook):
bq_job_id = "foo.bq.job_id"
mock_hook.return_value \
.run_query.return_value = bq_job_id
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_query.return_value = bq_job_id
mock_client.return_value \
.get_job.side_effects = [Exception("bq error")]
# To make sure hasattr "sees" close and calls it
mock_client.return_value.close.return_value
mock.seal(mock_hook)
execution_date = datetime.utcnow().replace(tzinfo=pytz.utc)
dag = DAG(dag_id='TestBigQueryExtractorE2E')
dag.create_dagrun(run_id=str(uuid.uuid4()),
state=State.QUEUED,
execution_date=execution_date)
task = BigQueryExecuteQueryOperator(
sql='select first_name, last_name from dataset.customers;',
task_id="task_id",
dag=dag,
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0),
do_xcom_push=False)
task_instance = TaskInstance(task=task, execution_date=execution_date)
bq_extractor = BigQueryExtractor(task)
tasks_meta_extract = bq_extractor.extract()
assert tasks_meta_extract is None
task_instance.run()
task_meta = bq_extractor.extract_on_complete(task_instance)
assert task_meta.run_facets['bigQuery_error'] == BigQueryErrorRunFacet(
clientError=mock.ANY)
mock_client.return_value.get_job.assert_called_once_with(
job_id=bq_job_id)
assert task_meta.inputs is not None
assert len(task_meta.inputs) == 0
assert task_meta.outputs is not None
assert len(task_meta.outputs) == 0
mock_client.return_value.close.assert_called()
def test_seals_recurse_on_added_attributes(self):
m = mock.Mock()
m.test1.test2().test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2().test3, 4)
with self.assertRaises(AttributeError):
m.test1.test2().test4
with self.assertRaises(AttributeError):
m.test1.test3
def test_seals_recurse_on_added_attributes(self):
m = mock.Mock()
m.test1.test2().test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2().test3, 4)
with self.assertRaises(AttributeError):
m.test1.test2().test4
with self.assertRaises(AttributeError):
m.test1.test3
def test_integration_with_spec_att_definition(self):
"""You are not restricted when using mock with spec"""
m = mock.Mock(SampleObject)
m.attr_sample1 = 1
m.attr_sample3 = 3
mock.seal(m)
self.assertEqual(m.attr_sample1, 1)
self.assertEqual(m.attr_sample3, 3)
with self.assertRaises(AttributeError):
m.attr_sample2
def test_integration_with_spec_att_definition(self):
"""You are not restricted when using mock with spec"""
m = mock.Mock(SampleObject)
m.attr_sample1 = 1
m.attr_sample3 = 3
mock.seal(m)
self.assertEqual(m.attr_sample1, 1)
self.assertEqual(m.attr_sample3, 3)
with self.assertRaises(AttributeError):
m.attr_sample2
def test_seals_recurse_on_magic_methods(self):
m = mock.MagicMock()
m.test1.test2["a"].test3 = 4
m.test1.test3[2:5].test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2["a"].test3, 4)
self.assertEqual(m.test1.test2[2:5].test3, 4)
with self.assertRaises(AttributeError):
m.test1.test2["a"].test4
with self.assertRaises(AttributeError):
m.test1.test3[2:5].test4
def test_seals_recurse_on_magic_methods(self):
m = mock.MagicMock()
m.test1.test2["a"].test3 = 4
m.test1.test3[2:5].test3 = 4
mock.seal(m)
self.assertEqual(m.test1.test2["a"].test3, 4)
self.assertEqual(m.test1.test2[2:5].test3, 4)
with self.assertRaises(AttributeError):
m.test1.test2["a"].test4
with self.assertRaises(AttributeError):
m.test1.test3[2:5].test4
def test_extract_error(self, mock_client, mock_hook):
bq_job_id = "foo.bq.job_id"
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_query.return_value = bq_job_id
mock_client.return_value \
.get_job.side_effects = [Exception("bq error")]
# To make sure hasattr "sees" close and calls it
mock_client.return_value.close.return_value
mock.seal(mock_hook)
mock.seal(mock_client)
dag = DAG(dag_id='TestBigQueryExtractorE2E')
task = BigQueryOperator(
sql='select first_name, last_name from customers;',
task_id="task_id",
project_id="project_id",
dag_id="dag_id",
dag=dag,
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0)
)
task_instance = TaskInstance(
task=task,
execution_date=datetime.utcnow().replace(tzinfo=pytz.utc))
bq_extractor = BigQueryExtractor(task)
steps_meta_extract = bq_extractor.extract()
assert steps_meta_extract is None
task_instance.run()
step_meta = bq_extractor.extract_on_complete(task_instance)
assert step_meta.context['bigquery.extractor.error'] is not None
mock_client.return_value \
.get_job.assert_called_once_with(job_id=bq_job_id)
assert step_meta.inputs is not None
assert len(step_meta.inputs) == 0
assert step_meta.outputs is not None
assert len(step_meta.outputs) == 0
assert step_meta.context['sql'] == task.sql
mock_client.return_value.close.assert_called()
def test_extract_cached(self, mock_client, mock_hook):
bq_job_id = "foo.bq.job_id"
mock_hook.return_value \
.run_query.return_value = bq_job_id
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_query.return_value = bq_job_id
job_details = self.read_file_json(
"tests/extractors/cached_job_details.json")
mock_client.return_value \
.get_job.return_value \
._properties = job_details
# To make sure hasattr "sees" close and calls it
mock_client.return_value.close.return_value
mock.seal(mock_client)
execution_date = datetime.utcnow().replace(tzinfo=pytz.utc)
dag = DAG(dag_id='TestBigQueryExtractorE2E')
dag.create_dagrun(run_id=str(uuid.uuid4()),
state=State.QUEUED,
execution_date=execution_date)
task = BigQueryExecuteQueryOperator(
sql='select first_name, last_name from dataset.customers;',
task_id="task_id",
dag=dag,
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0),
do_xcom_push=False)
task_instance = TaskInstance(task=task, execution_date=execution_date)
bq_extractor = BigQueryExtractor(task)
tasks_meta_extract = bq_extractor.extract()
assert tasks_meta_extract is None
task_instance.run()
task_meta = bq_extractor.extract_on_complete(task_instance)
assert task_meta.inputs is not None
assert task_meta.outputs is not None
assert len(task_meta.run_facets) == 1
assert task_meta.run_facets['bigQuery_job'] \
== BigQueryJobRunFacet(cached=True)
def test_extract_cached(self, mock_client, mock_hook):
bq_job_id = "foo.bq.job_id"
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_query.return_value = bq_job_id
job_details = self.read_file_json(
"tests/extractors/cached_job_details.json"
)
mock_client.return_value.get_job.return_value._properties = job_details
# To make sure hasattr "sees" close and calls it
mock_client.return_value.close.return_value
mock.seal(mock_hook)
mock.seal(mock_client)
dag = DAG(dag_id='TestBigQueryExtractorE2E')
task = BigQueryOperator(
sql='select first_name, last_name from customers;',
task_id="task_id",
project_id="project_id",
dag_id="dag_id",
dag=dag,
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0)
)
task_instance = TaskInstance(
task=task,
execution_date=datetime.utcnow().replace(tzinfo=pytz.utc))
bq_extractor = BigQueryExtractor(task)
steps_meta_extract = bq_extractor.extract()
assert steps_meta_extract is None
task_instance.run()
step_meta = bq_extractor.extract_on_complete(task_instance)
assert step_meta.inputs is not None
assert step_meta.outputs is not None
assert len(step_meta.run_facets) == 1
assert step_meta.run_facets['bigQuery_statistics'] \
== BigQueryStaticticsRunFacet(cached=True)
def test_extract(self, mock_client, mock_hook):
log.info("test_extractor")
job_details = self.read_file_json(
"tests/extractors/job_details.json")
table_details = self.read_dataset_json(
"tests/extractors/table_details.json")
out_details = self.read_dataset_json(
"tests/extractors/out_table_details.json")
bq_job_id = "foo.bq.job_id"
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_query.return_value = bq_job_id
mock_client.return_value \
.get_job.return_value \
._properties = job_details
mock_client.return_value \
.get_table.side_effect = [table_details, out_details]
# To make sure hasattr "sees" close and calls it
mock_client.return_value.close.return_value
mock.seal(mock_hook)
mock.seal(mock_client)
dag = DAG(dag_id='TestBigQueryExtractorE2E')
task = BigQueryOperator(
sql='select first_name, last_name from customers;',
task_id="task_id",
project_id="project_id",
dag_id="dag_id",
dag=dag,
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0)
)
task_instance = TaskInstance(
task=task,
execution_date=datetime.utcnow().replace(tzinfo=pytz.utc))
bq_extractor = BigQueryExtractor(task)
steps_meta_extract = bq_extractor.extract()
assert steps_meta_extract is None
task_instance.run()
step_meta = bq_extractor.extract_on_complete(task_instance)
assert step_meta.context['bigquery.job_properties'] \
== json.dumps(job_details)
mock_client.return_value \
.get_job.assert_called_once_with(job_id=bq_job_id)
assert step_meta.inputs is not None
assert len(step_meta.inputs) == 1
assert step_meta.inputs[0].name == \
'bigquery-public-data.usa_names.usa_1910_2013'
assert step_meta.inputs[0].fields is not None
assert len(step_meta.inputs[0].fields) == 5
assert step_meta.outputs is not None
assert len(step_meta.outputs) == 1
assert step_meta.outputs[0].fields is not None
assert len(step_meta.outputs[0].fields) == 2
assert step_meta.outputs[0].name == \
'bq-airflow-marquez.new_dataset.output_table'
assert step_meta.context['sql'] == task.sql
assert step_meta.context['bigquery.job_id'] == bq_job_id
mock_client.return_value.close.assert_called()
def test_call_on_defined_sealed_mock_succeeds(self):
m = mock.Mock(return_value=5)
mock.seal(m)
self.assertEqual(m(), 5)
def test_call_on_sealed_mock_fails(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError):
m()
def test_initialized_attributes_allowed_after_seal(self):
m = mock.Mock(test_value=1)
mock.seal(m)
self.assertEqual(m.test_value, 1)
def test_initialized_attributes_allowed_after_seal(self):
m = mock.Mock(test_value=1)
mock.seal(m)
self.assertEqual(m.test_value, 1)
def test_call_on_sealed_mock_fails(self):
m = mock.Mock()
mock.seal(m)
with self.assertRaises(AttributeError):
m()
def test_call_on_defined_sealed_mock_succeeds(self):
m = mock.Mock(return_value=5)
mock.seal(m)
self.assertEqual(m(), 5)
