diff --git a/tensorboard/uploader/uploader_test.py b/tensorboard/uploader/uploader_test.py index 7f42cf9c86..3b926736cf 100644 --- a/tensorboard/uploader/uploader_test.py +++ b/tensorboard/uploader/uploader_test.py @@ -68,11 +68,47 @@ def _create_mock_client(): return mock_client +# Sentinel for `_create_*` helpers, for arguments for which we want to +# supply a default other than the `None` used by the code under test. +_USE_DEFAULT = object() + + +def _create_uploader( + writer_client=_USE_DEFAULT, + logdir=None, + rpc_rate_limiter=_USE_DEFAULT, + name=None, + description=None, +): + if writer_client is _USE_DEFAULT: + writer_client = _create_mock_client() + if rpc_rate_limiter is _USE_DEFAULT: + rpc_rate_limiter = util.RateLimiter(0) + return uploader_lib.TensorBoardUploader( + writer_client, + logdir, + rpc_rate_limiter=rpc_rate_limiter, + name=name, + description=description, + ) + + +def _create_request_sender( + experiment_id=None, api=None, rpc_rate_limiter=_USE_DEFAULT +): + if api is _USE_DEFAULT: + api = _create_mock_client() + if rpc_rate_limiter is _USE_DEFAULT: + rpc_rate_limiter = util.RateLimiter(0) + return uploader_lib._BatchedRequestSender( + experiment_id=experiment_id, api=api, rpc_rate_limiter=rpc_rate_limiter + ) + + class TensorboardUploaderTest(tf.test.TestCase): def test_create_experiment(self): logdir = "/logs/foo" - mock_client = _create_mock_client() - uploader = uploader_lib.TensorBoardUploader(mock_client, logdir) + uploader = _create_uploader(_create_mock_client(), logdir) eid = uploader.create_experiment() self.assertEqual(eid, "123") @@ -80,9 +116,7 @@ def test_create_experiment_with_name(self): logdir = "/logs/foo" mock_client = _create_mock_client() new_name = "This is the new name" - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, name=new_name - ) + uploader = _create_uploader(mock_client, logdir, name=new_name) eid = uploader.create_experiment() self.assertEqual(eid, "123") mock_client.CreateExperiment.assert_called_once() @@ -100,7 +134,7 @@ def test_create_experiment_with_description(self): **description**" may have "strange" unicode chars 🌴 \\/<> """ - uploader = uploader_lib.TensorBoardUploader( + uploader = _create_uploader( mock_client, logdir, description=new_description ) eid = uploader.create_experiment() @@ -121,7 +155,7 @@ def test_create_experiment_with_all_metadata(self): may have "strange" unicode chars 🌴 \/<> """ new_name = "This is a cool name." - uploader = uploader_lib.TensorBoardUploader( + uploader = _create_uploader( mock_client, logdir, name=new_name, description=new_description ) eid = uploader.create_experiment() @@ -136,7 +170,7 @@ def test_create_experiment_with_all_metadata(self): def test_start_uploading_without_create_experiment_fails(self): mock_client = _create_mock_client() - uploader = uploader_lib.TensorBoardUploader(mock_client, "/logs/foo") + uploader = _create_uploader(mock_client, "/logs/foo") with self.assertRaisesRegex(RuntimeError, "call create_experiment()"): uploader.start_uploading() @@ -145,8 +179,8 @@ def test_start_uploading_without_create_experiment_fails(self): def test_start_uploading(self): mock_client = _create_mock_client() mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, "/logs/foo", mock_rate_limiter + uploader = _create_uploader( + mock_client, "/logs/foo", rpc_rate_limiter=mock_rate_limiter ) uploader.create_experiment() @@ -177,10 +211,7 @@ def scalar_event(tag, value): def test_upload_empty_logdir(self): logdir = self.get_temp_dir() mock_client = _create_mock_client() - mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, mock_rate_limiter - ) + uploader = _create_uploader(mock_client, logdir) uploader.create_experiment() uploader._upload_once() mock_client.WriteScalar.assert_not_called() @@ -190,10 +221,7 @@ def test_upload_swallows_rpc_failure(self): with tb_test_util.FileWriter(logdir) as writer: writer.add_test_summary("foo") mock_client = _create_mock_client() - mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, mock_rate_limiter - ) + uploader = _create_uploader(mock_client, logdir) uploader.create_experiment() error = test_util.grpc_error(grpc.StatusCode.INTERNAL, "Failure") mock_client.WriteScalar.side_effect = error @@ -205,10 +233,7 @@ def test_upload_propagates_experiment_deletion(self): with tb_test_util.FileWriter(logdir) as writer: writer.add_test_summary("foo") mock_client = _create_mock_client() - mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, mock_rate_limiter - ) + uploader = _create_uploader(mock_client, logdir) uploader.create_experiment() error = test_util.grpc_error(grpc.StatusCode.NOT_FOUND, "nope") mock_client.WriteScalar.side_effect = error @@ -227,10 +252,7 @@ def test_upload_preserves_wall_time(self): ) ) mock_client = _create_mock_client() - mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, mock_rate_limiter - ) + uploader = _create_uploader(mock_client, logdir) uploader.create_experiment() uploader._upload_once() mock_client.WriteScalar.assert_called_once() @@ -245,10 +267,7 @@ def test_upload_preserves_wall_time(self): def test_upload_full_logdir(self): logdir = self.get_temp_dir() mock_client = _create_mock_client() - mock_rate_limiter = mock.create_autospec(util.RateLimiter) - uploader = uploader_lib.TensorBoardUploader( - mock_client, logdir, mock_rate_limiter - ) + uploader = _create_uploader(mock_client, logdir) uploader.create_experiment() # Convenience helpers for constructing expected requests. @@ -359,11 +378,7 @@ def test_upload_full_logdir(self): class BatchedRequestSenderTest(tf.test.TestCase): def _populate_run_from_events(self, run_proto, events): mock_client = _create_mock_client() - builder = uploader_lib._BatchedRequestSender( - experiment_id="123", - api=mock_client, - rpc_rate_limiter=util.RateLimiter(0), - ) + builder = _create_request_sender(experiment_id="123", api=mock_client) builder.send_requests({"": events}) requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list] if requests: @@ -549,9 +564,7 @@ def test_no_budget_for_experiment_id(self): long_experiment_id = "A" * uploader_lib._MAX_REQUEST_LENGTH_BYTES mock_client = _create_mock_client() with self.assertRaises(RuntimeError) as cm: - builder = uploader_lib._BatchedRequestSender( - long_experiment_id, mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender(long_experiment_id, mock_client) builder.send_requests(run_to_events) self.assertEqual( str(cm.exception), "Byte budget too small for experiment ID" @@ -564,9 +577,7 @@ def test_no_room_for_single_point(self): long_run_name = "A" * uploader_lib._MAX_REQUEST_LENGTH_BYTES run_to_events = {long_run_name: [event]} with self.assertRaises(RuntimeError) as cm: - builder = uploader_lib._BatchedRequestSender( - "123", mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender("123", mock_client) builder.send_requests(run_to_events) self.assertEqual(str(cm.exception), "add_event failed despite flush") @@ -584,9 +595,7 @@ def test_break_at_run_boundary(self): [(long_run_1, [event_1]), (long_run_2, [event_2])] ) - builder = uploader_lib._BatchedRequestSender( - "123", mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender("123", mock_client) builder.send_requests(run_to_events) requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list] @@ -623,9 +632,7 @@ def test_break_at_tag_boundary(self): event.summary.value.add(tag=long_tag_2, simple_value=2.0) run_to_events = {"train": [event]} - builder = uploader_lib._BatchedRequestSender( - "123", mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender("123", mock_client) builder.send_requests(run_to_events) requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list] for request in requests: @@ -665,9 +672,7 @@ def test_break_at_scalar_point_boundary(self): events.append(event_pb2.Event(summary=summary, step=step)) run_to_events = {"train": events} - builder = uploader_lib._BatchedRequestSender( - "123", mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender("123", mock_client) builder.send_requests(run_to_events) requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list] for request in requests: @@ -722,9 +727,7 @@ def mock_create_point(uploader_self, *args, **kwargs): "_create_point", mock_create_point, ): - builder = uploader_lib._BatchedRequestSender( - "123", mock_client, util.RateLimiter(0) - ) + builder = _create_request_sender("123", mock_client) builder.send_requests(run_to_events) requests = [c[0][0] for c in mock_client.WriteScalar.call_args_list] for request in requests: