def patch_sim_scene(module_path: str, **kwargs):
"""Patches the SimScene class in the given module.
Args:
module_path: The path to the SimScene class to mock.
**kwargs: Arguments passed to MockSimScene when `SimScene.create` is
called.
"""
with mock.patch(module_path, MockSimScene) as mock_sim_cls:
mock_sim_cls.create = lambda *_args, **_kwargs: MockSimScene(**kwargs)
yield mock_sim_cls
def test_ga_hook_send_batch_hit_with_return_error_status_code(self, r_code):
"""Test GoogleAnalyticsHook sends batch hit with error response."""
test_hook = ga_hook.GoogleAnalyticsHook(self.test_tracking_id,
self.event_test_data,
False)
test_payload = self.payload_builder.generate_batch_payload(
'event', [self.event_test_data])
with mock.patch('urllib.request.urlopen') as urlopen_mock:
mock_response = urlopen_mock.return_value.__enter__.return_value
mock_response.status = r_code
with self.assertRaises(errors.DataOutConnectorSendUnsuccessfulError):
test_hook.send_hit(test_payload, send_type='batch')
def test_ga_hook_send_single_hit(self):
"""Test GoogleAnalyticsHook sends single hit."""
test_hook = ga_hook.GoogleAnalyticsHook(self.test_tracking_id,
self.event_test_data,
False)
test_payload = self.payload_builder.generate_single_payload(
'event', self.event_test_data)
with mock.patch('urllib.request.urlopen') as urlopen_mock:
mock_response = urlopen_mock.return_value.__enter__.return_value
mock_response.status = 200
test_hook.send_hit(test_payload)
urlopen_mock.assert_called_once()
def test_ga_hook_send_batch_hit_with_retrys_on_retriable_error(self):
"""Test GoogleAnalyticsHook retries when retriable error occures."""
test_hook = ga_hook.GoogleAnalyticsHook(self.test_tracking_id,
self.event_test_data,
False)
test_payload = self.payload_builder.generate_batch_payload(
'event', [self.event_test_data])
with mock.patch('urllib.request.urlopen') as urlopen_mock:
urlopen_mock.return_value.__enter__.return_value.status = 429
try:
test_hook.send_hit(test_payload)
except errors.DataOutConnectorSendUnsuccessfulError:
pass
self.assertEqual(urlopen_mock.call_count,
retry_utils._RETRY_UTILS_MAX_RETRIES)
def setUp(self):
super().setUp()
# Set a constant version for artifact version tag.
patcher = mock.patch('tfx.version.__version__')
patcher.start()
tfx_version.__version__ = '0.123.4.dev'
self.addCleanup(patcher.stop)
pipeline_root = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self.id())
metadata_path = os.path.join(pipeline_root, 'metadata', 'metadata.db')
self._metadata_path = metadata_path
connection_config = metadata.sqlite_metadata_connection_config(
metadata_path)
connection_config.sqlite.SetInParent()
self._mlmd_connection = metadata.Metadata(
connection_config=connection_config)
pipeline = self._make_pipeline(pipeline_root, str(uuid.uuid4()))
self._pipeline = pipeline
self._importer_node = self._pipeline.nodes[0].pipeline_node
self._task_queue = tq.TaskQueue()
[importer_task] = test_utils.run_generator_and_test(
test_case=self,
mlmd_connection=self._mlmd_connection,
generator_class=sptg.SyncPipelineTaskGenerator,
pipeline=self._pipeline,
task_queue=self._task_queue,
use_task_queue=True,
service_job_manager=None,
num_initial_executions=0,
num_tasks_generated=1,
num_new_executions=1,
num_active_executions=1,
expected_exec_nodes=[self._importer_node])
self._importer_task = importer_task
def patch_time(module_path: str, **kwargs):
return mock.patch(module_path, MockTime(**kwargs))
def patched_fn(self):
DynamixelRobotComponent.DEVICE_CLIENTS.clear()
with mock.patch(
'robel.components.robot.dynamixel_robot.DynamixelClient',
new=MockDynamixelClient):
test_fn(self)
