Skip to content

Commit

Permalink
feat: set primary user agent
Browse files Browse the repository at this point in the history
Signed-off-by: Michal Fiedorowicz <[email protected]>
  • Loading branch information
mfiedorowicz committed Sep 9, 2024
1 parent d0a9729 commit 3a7e147
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 0 deletions.
6 changes: 6 additions & 0 deletions netboxlabs/diode/sdk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,18 +102,24 @@ def __init__(
("python-version", self._python_version),
)

channel_opts = (
("grpc.primary_user_agent", f"{self._name}/{self._version} {self._app_name}/{self._app_version}"),
)

if self._tls_verify:
_LOGGER.debug("Setting up gRPC secure channel")
self._channel = grpc.secure_channel(
self._target,
grpc.ssl_channel_credentials(
root_certificates=_load_certs(),
),
options=channel_opts,
)
else:
_LOGGER.debug("Setting up gRPC insecure channel")
self._channel = grpc.insecure_channel(
target=self._target,
options=channel_opts,
)

channel = self._channel
Expand Down
40 changes: 40 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,46 @@ def test_client_sets_up_insecure_channel_when_grpc_scheme_is_found_in_target():
mock_insecure_channel.assert_called_once()


def test_insecure_channel_options_with_primary_user_agent():
"""Check that DiodeClient.__init__() sets the gRPC primary_user_agent option for insecure channel."""
with mock.patch("grpc.insecure_channel") as mock_insecure_channel:
client = DiodeClient(
target="grpc://localhost:8081",
app_name="my-producer",
app_version="0.0.1",
api_key="abcde",
)

mock_insecure_channel.assert_called_once()
_, kwargs = mock_insecure_channel.call_args
assert kwargs["options"] == (
(
"grpc.primary_user_agent",
f"{client.name}/{client.version} {client.app_name}/{client.app_version}",
),
)


def test_secure_channel_options_with_primary_user_agent():
"""Check that DiodeClient.__init__() sets the gRPC primary_user_agent option for secure channel."""
with mock.patch("grpc.secure_channel") as mock_secure_channel:
client = DiodeClient(
target="grpcs://localhost:8081",
app_name="my-producer",
app_version="0.0.1",
api_key="abcde",
)

mock_secure_channel.assert_called_once()
_, kwargs = mock_secure_channel.call_args
assert kwargs["options"] == (
(
"grpc.primary_user_agent",
f"{client.name}/{client.version} {client.app_name}/{client.app_version}",
),
)


def test_client_interceptor_setup_with_path():
"""Check that DiodeClient.__init__() sets up the gRPC interceptor when a path is provided."""
client = DiodeClient(
Expand Down

0 comments on commit 3a7e147

Please sign in to comment.