From fc8b0f92c88feac7f62d8c98aa13e8fde75ea157 Mon Sep 17 00:00:00 2001
From: Michal Fiedorowicz <mfiedorowicz@netboxlabs.com>
Date: Mon, 9 Sep 2024 11:29:18 +0100
Subject: [PATCH] feat: set primary user agent (#15)

* feat: set primary user agent

Signed-off-by: Michal Fiedorowicz <mfiedorowicz@netboxlabs.com>

---------

Signed-off-by: Michal Fiedorowicz <mfiedorowicz@netboxlabs.com>
---
 .github/CODEOWNERS             |  2 +-
 netboxlabs/diode/sdk/client.py |  6 +++++
 tests/test_client.py           | 40 ++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 9c13251..47b6708 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1 +1 @@
-*                     @leoparente @mfiedorowicz @natm
+*                     @leoparente @mfiedorowicz
diff --git a/netboxlabs/diode/sdk/client.py b/netboxlabs/diode/sdk/client.py
index 09ee255..c62bae6 100644
--- a/netboxlabs/diode/sdk/client.py
+++ b/netboxlabs/diode/sdk/client.py
@@ -102,6 +102,10 @@ 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(
@@ -109,11 +113,13 @@ def __init__(
                 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
diff --git a/tests/test_client.py b/tests/test_client.py
index cbf0f52..f1d6596 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -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(