def test_read_raises_blocking_error(self):
file_object = mock.Mock()
file_object.read = mock.Mock(
side_effect=io.BlockingIOError(None, None))
adapter = file_object_adapter.FileObjectAdapter(file_object)
self.assertEqual(adapter.read(10), b'')
def setUp(self):
"""Setup for CloudApiTest."""
super(CloudApiTest, self).setUp()
self.addCleanup(mock.patch.stopall)
self.project_id = 'project_id'
self.service_account_name = 'service_account_name'
self.api_name = 'storage'
self.session = mock.Mock()
self.mock_client = mock.Mock()
self.parent = f'projects/{self.project_id}'
self.service_account_key_file = '/path/to/service/account/key'
self.mock_credentials = mock.Mock(credentials.Credentials,
autospec=True)
# Mocks for building client.
mock_get_credentials = mock.patch.object(cloud_auth,
'get_credentials',
autospec=True).start()
mock_build_service_client = mock.patch.object(cloud_auth,
'build_service_client',
autospec=True).start()
mock_get_credentials.return_value = self.mock_credentials
mock_build_service_client.return_value = self.mock_client
# Mocks for operations.
self.operation_client = mock.Mock()
self.mock_wait_for_operation = mock.patch.object(
utils, 'wait_for_operation', autospec=True).start()
self.mock_execute_request = mock.patch.object(utils,
'execute_request',
autospec=True).start()
self.operation = {}
self.mock_execute_request.return_value = self.operation
self.mock_client.operations.return_value = self.operation_client
self.cloud_api_utils = cloud_api.CloudApiUtils(
self.project_id,
service_account_key_file=self.service_account_key_file)
def setUp(self):
"""Creates mock objects for googleapi client."""
super(CloudSchedulerTest, self).setUp()
self.addCleanup(mock.patch.stopall)
self.project_id = 'project_id'
self.location = 'us-central1'
self.service_account_key_file = '/tmp/service_account_key.json'
self.mock_get_credentials = mock.patch.object(cloud_auth,
'get_credentials',
autospec=True).start()
self.mock_build_service_client = mock.patch.object(
cloud_auth, 'build_service_client', autospec=True).start()
self.mock_client = mock.Mock()
self.mock_credentials = mock.Mock(credentials.Credentials,
autospec=True)
self.mock_get_credentials.return_value = self.mock_credentials
self.mock_build_service_client.return_value = self.mock_client
self.fake_appengine_http_target = cloud_scheduler.AppEngineTarget(
http_method='GET', relative_uri='/test', service='test')
self.fake_http_target = cloud_scheduler.HttpTarget(
http_method='POST',
uri='https://www.google.com/',
body='{}',
headers={'Content-Type': 'application/json'},
authorization_header=('*****@*****.**',
'my-fake-scope'))
self.scheduler = cloud_scheduler.CloudSchedulerUtils(
project_id=self.project_id,
location=self.location,
service_account_key_file=self.service_account_key_file)
def test_read_no_data_available(self):
f = mock.Mock()
f.read = mock.Mock(return_value=None)
f.readable = mock.Mock(return_value=True)
ds = get_decrypting_stream(f, B_AAD_)
self.assertRaises(io.BlockingIOError, ds.read, 5)
def test_write_raises_blocking_error(self):
file_object = mock.Mock()
file_object.write = mock.Mock(
side_effect=io.BlockingIOError(None, None, 5))
adapter = file_object_adapter.FileObjectAdapter(file_object)
self.assertEqual(5, adapter.write(b'something'))
def test_read_eof(self):
file_object = mock.Mock()
file_object.read = mock.Mock(return_value=b'')
adapter = file_object_adapter.FileObjectAdapter(file_object)
with self.assertRaises(EOFError):
adapter.read(10)
def test_reset_dict(self):
"""Performs a reset using dictionary observations."""
test = TestEnv(use_dict_obs=True)
test._reset = mock.Mock()
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([('a', [1, 1])]))
obs = test.reset()
self.assertEqual(test._reset.call_count, 1)
self.assertEqual(test.get_obs_dict.call_count, 1)
np.testing.assert_array_equal(obs['a'], [1, 1])
def test_last_action(self):
"""Tests reading the last action."""
test = TestEnv(nq=3)
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([('o', [0])]))
test.get_reward_dict = mock.Mock(return_value={'r': np.array(0)})
self.assertIsNone(test.last_action)
np.testing.assert_array_equal(test._get_last_action(), [0, 0, 0])
test.step([1, 1, 1])
np.testing.assert_array_equal(test._get_last_action(), [1, 1, 1])
self.assertIs(test.last_action, test._get_last_action())
def test_sticky_actions(self):
"""Tests sticky actions."""
test = TestEnv(nq=3, sticky_action_probability=1)
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([('o', [0])]))
test.get_reward_dict = mock.Mock(return_value={'r': np.array(0)})
test._step = mock.Mock()
test.last_action = np.ones(3)
test.step([0, 0, 0])
args, _ = test._step.call_args_list[0]
np.testing.assert_array_equal(args[0], np.ones(3))
def setUp(self):
super().setUp()
# Mock out constructor to avoid Exceptions reading from Airflow env
automl_tables_hook.AutoMLTablesHook.__init__ = mock.Mock(
return_value=None)
self.mock_predict = mock.Mock()
automl_tables_hook.AutoMLTablesHook.batch_predict = self.mock_predict
self.mock_input_path = mock.Mock(return_value='gs://input')
(automl_pred_op.AutoMLTablesBatchPredictionOperator._get_input_path
) = self.mock_input_path
self.mock_output_path = mock.Mock()
(automl_pred_op.AutoMLTablesBatchPredictionOperator._set_output_path
) = self.mock_output_path
def setUp(self):
"""Set up mocks for hook test."""
super().setUp()
# Mock constructor to avoid super class init Exceptions
automl_tables_hook.AutoMLTablesHook.__init__ = mock.Mock(
return_value=None)
# Mock project id
automl_tables_hook.AutoMLTablesHook.project_id = mock.PropertyMock(
return_value='test_project')
# Mock session
mock_session = mock.Mock()
automl_tables_hook.AutoMLTablesHook._get_authorized_session = mock_session
self.mock_get = mock_session.return_value.get
self.mock_post = mock_session.return_value.post
def get_tabledata(self,
dataset_id,
table_id,
max_results,
page_token=None,
start_index=None,
selected_fields=None):
"""Mock method of BigQueryCursor.get_tabledata()."""
self.dataset_id = dataset_id
self.table_id = table_id
self.max_results = max_results
self.page_token = page_token
self.start_index = start_index
self.selected_fields = selected_fields
if table_id.startswith('error') and start_index == 1:
next(self.data_generator)
response = mock.Mock()
response.reason = 'test_reason'
raise googleapiclient_errors.HttpError(resp=response,
content=b'test')
try:
return next(self.data_generator)
except StopIteration:
return None
def test_stop_node_services_exception_handling(self):
self._mock_service_job_manager.stop_node_services.side_effect = RuntimeError(
'test error')
self.assertFalse(self._wrapper.stop_node_services(
mock.Mock(), 'node2'))
self._mock_service_job_manager.stop_node_services.assert_called_once_with(
mock.ANY, 'node2')
def testStepHybrid(self):
obs_spec = array_spec.BoundedArraySpec((2, 3), np.int32, -10, 10)
action_spec = {
'discrete':
array_spec.BoundedArraySpec((1,), np.int32, 1, 3),
'continuous':
array_spec.ArraySpec((2,), np.float32)
}
mock_env = mock.Mock(
wraps=random_py_environment.RandomPyEnvironment(obs_spec, action_spec))
one_hot_action_wrapper = wrappers.OneHotActionWrapper(mock_env)
one_hot_action_wrapper.reset()
action = {
'discrete':
np.array([[0, 1, 0]]).astype(np.int32),
'continuous':
np.array([0.5, 0.3]).astype(np.float32)
}
one_hot_action_wrapper.step(action)
self.assertTrue(mock_env.step.called)
expected_action = {
'discrete':
np.array([2]),
'continuous':
np.array([0.5, 0.3])
}
np.testing.assert_array_almost_equal(
expected_action['discrete'], mock_env.step.call_args[0][0]['discrete'])
np.testing.assert_array_almost_equal(
expected_action['continuous'],
mock_env.step.call_args[0][0]['continuous'])
def setUp(self):
"""Setup function for each unit test."""
super(GoogleAnalyticsHookTest, self).setUp()
self.test_tracking_id = 'UA-12323-4'
self.payload_builder = ga_hook.PayloadBuilder(self.test_tracking_id)
self.event_test_data = {
'ec': 'ClientID',
'ea': 'test_event_action',
'el': '20190423',
'ev': 1,
'cid': '12345.456789'
}
self.small_event = {
'cid': '12345.67890',
'ec': 'ClientID',
'ea': 'test_event_action',
'el': '20190423',
'ev': 1,
'z': '1558517072202080'
}
# Both of the below are approx 4K of data
self.medium_event = {**self.small_event, 'ea': 'x' * 3800}
self.utf8_event = {**self.small_event, 'ea': b'\xf0\xa9\xb8\xbd' * 320}
self.test_hook = ga_hook.GoogleAnalyticsHook(self.test_tracking_id,
self.event_test_data)
self.test_hook._send_http_request = mock.MagicMock(autospec=True)
self.test_hook._send_http_request.return_value = mock.Mock(ok=True)
self.test_hook._send_http_request.return_value.status = 200
def testShuffle(self):
"""Test that dataset is being shuffled when asked."""
# Reward of 1 is given if action == (context % 3)
context = tf.reshape(tf.range(128), shape=[128, 1])
labels = tf.math.mod(context, 3)
batch_size = 32
dataset = (tf.data.Dataset.from_tensor_slices(
(context, labels)).repeat().shuffle(4 * batch_size))
reward_distribution = deterministic_reward_distribution(tf.eye(3))
# Note - shuffle should hapen *first* in call chain, so this
# test will fail if shuffle is called e.g. after batch or prefetch.
dataset.shuffle = mock.Mock(spec=dataset.shuffle,
side_effect=dataset.shuffle)
ce.ClassificationBanditEnvironment(dataset, reward_distribution,
batch_size)
dataset.shuffle.assert_not_called()
ce.ClassificationBanditEnvironment(dataset,
reward_distribution,
batch_size,
shuffle_buffer_size=3,
seed=7)
dataset.shuffle.assert_called_with(buffer_size=3,
reshuffle_each_iteration=True,
seed=7)
def testWithAdvantageFn(self, with_value_network):
advantage_fn = mock.Mock(
side_effect=lambda returns, _: returns)
value_network = (DummyValueNet(self._obs_spec) if with_value_network
else None)
agent = reinforce_agent.ReinforceAgent(
self._time_step_spec,
self._action_spec,
actor_network=DummyActorNet(
self._obs_spec, self._action_spec, unbounded_actions=False),
value_network=value_network,
advantage_fn=advantage_fn,
optimizer=None,
)
step_type = tf.constant(
[[ts.StepType.FIRST, ts.StepType.LAST, ts.StepType.FIRST,
ts.StepType.LAST]])
reward = tf.constant([[0, 0, 0, 0]], dtype=tf.float32)
discount = tf.constant([[1, 1, 1, 1]], dtype=tf.float32)
observations = tf.constant(
[[[1, 2], [1, 2], [1, 2], [1, 2]]], dtype=tf.float32)
time_steps = ts.TimeStep(step_type, reward, discount, observations)
actions = tf.constant([[[0], [1], [2], [3]]], dtype=tf.float32)
agent.total_loss(time_steps, actions, time_steps.reward, None)
advantage_fn.assert_called_once()
def __init__(self, *args, step_duration: float = 1, **kwargs):
"""Initializes a new mock SimScene."""
self.sim = mock.Mock()
self.sim.model = MockMjModel(*args, **kwargs)
self.model = self.sim.model
self.sim.data = self.model
self.data = self.model.data
self.step_duration = step_duration
self.close = mock.Mock()
self.advance = mock.Mock()
self.renderer = mock.Mock()
self.renderer.render_offscreen = lambda w, h, **_: np.zeros((w, h))
def setUp(self):
super(UtilsTest, self).setUp()
self.addCleanup(mock.patch.stopall)
self.operation_not_completed = {
'name': 'projects/my-proj/locations/us-central1/operations/'
'11111111-aaaa-2222-bbbb-111111111111',
'done': False,
'error': None
}
self.operation_completed = {
'name': 'projects/my-proj/locations/us-central1/operations/'
'22222222-aaaa-2222-bbbb-111111111111',
'done': True,
'error': None
}
self.operation_failed = {
'name': 'projects/my-proj/locations/us-central1/operations/'
'33333333-aaaa-2222-bbbb-111111111111',
'done': False,
'error': {
'message': 'Timeout occurred!'
}
}
self.sleep_mock = mock.patch.object(time, 'sleep',
autospec=True).start()
self.mock_client = mock.Mock()
def test_batch_predict_retry(self):
"""Test for a retried prediction.
Test for a batch prediction which fails for the first time, and after
retry succeeds.
"""
# Predict response
self.mock_post.side_effect = [
mock.Mock(status_code=400, text='Bad Request'),
mock.Mock(status_code=200,
json=mock.Mock(return_value={
'name': 'test_operation_name',
'metadata': {}
})),
]
# Query status response
self.mock_get.return_value.json.side_effect = [
{
'done': False,
'metadata': {}
},
{
'done': False,
'metadata': {}
},
{
'done': True,
'metadata': {
'batchPredictDetails': {
'outputInfo': {
'gcsOutputDirectory': 'gs://outputfile'
}
}
}
},
]
self.mock_get.return_value.status_code = 200
hook = automl_tables_hook.AutoMLTablesHook()
# In unit tests we do not want to sleep between polling
pred_result = hook.batch_predict(input_path='gs://input',
output_path='gs://output',
model_id='test_model_id',
compute_region='test_region',
poll_wait_time=0)
self.assertEqual(pred_result['done'], True)
def _task_manager(self, task_queue):
with tm.TaskManager(
mock.Mock(),
task_queue,
max_active_task_schedulers=1000,
max_dequeue_wait_secs=0.1,
process_all_queued_tasks_before_exit=True) as task_manager:
yield task_manager
def test_calls_forwarded_to_underlying_instance(self):
self.assertEqual(
service_jobs.ServiceStatus.SUCCESS,
self._wrapper.ensure_node_services(mock.Mock(), 'node1'))
self.assertTrue(self._wrapper.stop_node_services(mock.Mock(), 'node2'))
self.assertTrue(
self._wrapper.is_pure_service_node(mock.Mock(), 'node3'))
self.assertFalse(
self._wrapper.is_mixed_service_node(mock.Mock(), 'node4'))
self._mock_service_job_manager.ensure_node_services.assert_called_once_with(
mock.ANY, 'node1')
self._mock_service_job_manager.stop_node_services.assert_called_once_with(
mock.ANY, 'node2')
self._mock_service_job_manager.is_pure_service_node.assert_called_once_with(
mock.ANY, 'node3')
self._mock_service_job_manager.is_mixed_service_node.assert_called_once_with(
mock.ANY, 'node4')
def test_ensure_node_services_cleanup_on_failure(self):
self._mock_service_job_manager.ensure_node_services.return_value = (
service_jobs.ServiceStatus.FAILED)
self.assertEqual(
service_jobs.ServiceStatus.FAILED,
self._wrapper.ensure_node_services(mock.Mock(), 'node1'))
self._mock_service_job_manager.stop_node_services.assert_called_once_with(
mock.ANY, 'node1')
def setUp(self):
super().setUp()
self.db = fake_firestore.FakeFirestore()
self.topic_path = 'MockTopicPath'
self.max_parallel_tasks = 3
self.mock_pubsub = mock.Mock()
self.mock_pubsub.topic_path.return_value = self.topic_path
def test_reset(self):
"""Performs a reset."""
test = TestEnv()
test._reset = mock.Mock()
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([('a', [1, 1])]))
obs = test.reset()
self.assertIsNone(test.last_action)
self.assertDictEqual(test.last_obs_dict, {'a': [1, 1]})
self.assertIsNone(test.last_reward_dict)
self.assertIsNone(test.last_score_dict)
self.assertFalse(test.is_done)
self.assertEqual(test.step_count, 0)
self.assertEqual(test._reset.call_count, 1)
np.testing.assert_array_equal(obs, [1, 1])
def test_scheduler_not_found(self):
task = test_utils.create_exec_node_task(node_uid=task_lib.NodeUid(
pipeline_uid=task_lib.PipelineUid(pipeline_id='pipeline'),
node_id='Transform'),
pipeline=self._pipeline)
with self.assertRaisesRegex(ValueError, 'No task scheduler found'):
ts.TaskSchedulerRegistry.create_task_scheduler(
mock.Mock(), self._pipeline, task)
def test_step(self):
"""Checks that step calls its subcomponents."""
test = TestEnv()
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([
('o1', [0]),
('o2', [1, 2]),
]))
test.get_reward_dict = mock.Mock(return_value={
'r1': np.array(1),
'r2': np.array(2),
})
test.get_score_dict = mock.Mock(return_value={
's1': np.array(1),
})
test.get_done = mock.Mock(return_value=np.array(False))
obs, reward, done, info = test.step([0])
np.testing.assert_array_equal(test.last_action, [0])
self.assertDictEqual(test.last_obs_dict, {'o1': [0], 'o2': [1, 2]})
self.assertDictEqual(test.last_reward_dict, {'r1': 1, 'r2': 2})
self.assertDictEqual(test.last_score_dict, {'s1': 1})
self.assertFalse(test.is_done)
self.assertEqual(test.step_count, 1)
# Check that step calls its sub-methods exactly once.
self.assertEqual(test.get_obs_dict.call_count, 1)
self.assertEqual(test.get_reward_dict.call_count, 1)
self.assertEqual(test.get_score_dict.call_count, 1)
self.assertEqual(test.get_done.call_count, 1)
np.testing.assert_array_equal(obs, [0, 1, 2])
self.assertEqual(reward, 3)
self.assertFalse(done)
self.assertListEqual(
sorted(info.keys()), [
'obs/o1', 'obs/o2', 'reward/r1', 'reward/r2', 'reward/total',
'score/s1'
])
np.testing.assert_array_equal(info['obs/o1'], [0])
np.testing.assert_array_equal(info['obs/o2'], [1, 2])
np.testing.assert_array_equal(info['reward/r1'], 1)
np.testing.assert_array_equal(info['reward/r2'], 2)
np.testing.assert_array_equal(info['reward/total'], 3)
np.testing.assert_array_equal(info['score/s1'], 1)
def test_step_dict(self):
"""Checks the output of step using dictionary observations."""
test = TestEnv(use_dict_obs=True)
test.get_obs_dict = mock.Mock(
return_value=collections.OrderedDict([
('o1', [0]),
('o2', [1, 2]),
]))
test.get_reward_dict = mock.Mock(return_value={
'r1': np.array(1),
'r2': np.array(2),
})
obs, _, _, _ = test.step([0])
self.assertListEqual(sorted(obs.keys()), ['o1', 'o2'])
np.testing.assert_array_equal(obs['o1'], [0])
np.testing.assert_array_equal(obs['o2'], [1, 2])
def test_execute_request_retries_on_service_unavailable_http_error(self):
mock_request = mock.Mock(http.HttpRequest)
content = b''
error = errors.HttpError(mock.MagicMock(status=503), content)
mock_request.execute.side_effect = [error, None]
utils.execute_request(mock_request)
self.assertEqual(mock_request.execute.call_count, 2)
def test_step(self):
action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 4)
env = _build_test_env(action_spec=action_spec)
mock_env = mock.Mock(wraps=env)
wrapper = tf_wrappers.OneHotActionWrapper(mock_env)
wrapper.reset()
wrapper.step(tf.constant([[0, 1, 0, 0, 0], [0, 0, 0, 1, 0]], tf.int32))
self.assertTrue(mock_env.step.called)
self.assertAllEqual([1, 3], mock_env.step.call_args[0][0])
