Skip to content

Commit

Permalink
uploader: set user agent for frontend handshake (#2934)
Browse files Browse the repository at this point in the history
Test Plan:
Running against a local frontend server, `tensorboard/2.1.0a0` shows up
in the server logs where previously there was `python-requests/2.22.0`.
Unit tests also included.

wchargin-branch: uploader-user-agent
  • Loading branch information
wchargin committed Nov 25, 2019
1 parent a5f5d46 commit 61da2a9
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
5 changes: 4 additions & 1 deletion tensorboard/uploader/server_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,10 @@ def fetch_server_info(origin):
post_body = _server_info_request().SerializeToString()
try:
response = requests.post(
endpoint, data=post_body, timeout=_REQUEST_TIMEOUT_SECONDS
endpoint,
data=post_body,
timeout=_REQUEST_TIMEOUT_SECONDS,
headers={"User-Agent": "tensorboard/%s" % version.VERSION},
)
except requests.RequestException as e:
raise CommunicationError("Failed to connect to backend: %s" % e)
Expand Down
12 changes: 12 additions & 0 deletions tensorboard/uploader/server_info_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,18 @@ def app(request):
self.assertIn("Corrupt response from backend", msg)
self.assertIn("an unlikely proto", msg)

def test_user_agent(self):
@wrappers.BaseRequest.application
def app(request):
result = server_info_pb2.ServerInfoResponse()
result.compatibility.details = request.headers["User-Agent"]
return wrappers.BaseResponse(result.SerializeToString())

origin = self._start_server(app)
result = server_info.fetch_server_info(origin)
expected_user_agent = "tensorboard/%s" % version.VERSION
self.assertEqual(result.compatibility.details, expected_user_agent)


class CreateServerInfoTest(tb_test.TestCase):
"""Tests for `create_server_info`."""
Expand Down

0 comments on commit 61da2a9

Please sign in to comment.