fix import

Author: leoparente

tests/test_client.py

@@ -22,7 +22,7 @@
     load_dryrun_entities,
     parse_target,
 )
-from netboxlabs.diode.sdk.diode.v1.ingester_pb2 import Entity as EntityPb
+from netboxlabs.diode.sdk.diode.v1 import ingester_pb2
 from netboxlabs.diode.sdk.exceptions import DiodeClientError, DiodeConfigError
 from netboxlabs.diode.sdk.ingester import Entity
 from netboxlabs.diode.sdk.version import version_semver
@@ -64,7 +64,9 @@ def test_config_errors(client_id, client_secret, env_var_name):
             client_id=client_id,
             client_secret=client_secret,
         )
-    assert str(err.value) == f"parameter or {env_var_name} environment variable required"
+    assert (
+        str(err.value) == f"parameter or {env_var_name} environment variable required"
+    )
 
 
 def test_client_error(mock_diode_authentication):
@@ -90,7 +92,10 @@ def test_diode_client_error_repr_returns_correct_string():
     error = DiodeClientError(grpc_error)
     error._status_code = grpc.StatusCode.UNAVAILABLE
     error._details = "Some details about the error"
-    assert repr(error) == "<DiodeClientError status code: StatusCode.UNAVAILABLE, details: Some details about the error>"
+    assert (
+        repr(error)
+        == "<DiodeClientError status code: StatusCode.UNAVAILABLE, details: Some details about the error>"
+    )
 
 
 def test_load_certs_returns_bytes():
@@ -180,7 +185,9 @@ def test_setup_sentry_initializes_with_correct_parameters(mock_diode_authenticat
         )
 
 
-def test_client_sets_up_secure_channel_when_grpcs_scheme_is_found_in_target(mock_diode_authentication):
+def test_client_sets_up_secure_channel_when_grpcs_scheme_is_found_in_target(
+    mock_diode_authentication,
+):
     """Check that DiodeClient.__init__() sets up the gRPC secure channel when grpcs:// scheme is found in the target."""
     client = DiodeClient(
         target="grpcs://localhost:8081",
@@ -205,7 +212,9 @@ def test_client_sets_up_secure_channel_when_grpcs_scheme_is_found_in_target(mock
         mock_secure_channel.assert_called_once()
 
 
-def test_client_sets_up_insecure_channel_when_grpc_scheme_is_found_in_target(mock_diode_authentication):
+def test_client_sets_up_insecure_channel_when_grpc_scheme_is_found_in_target(
+    mock_diode_authentication,
+):
     """Check that DiodeClient.__init__() sets up the gRPC insecure channel when grpc:// scheme is found in the target."""
     client = DiodeClient(
         target="grpc://localhost:8081",
@@ -347,10 +356,14 @@ def test_client_setup_sentry_called_when_sentry_dsn_exists(mock_diode_authentica
             client_secret="123456",
             sentry_dsn="https://user@password.mock.dsn/123456",
         )
-        mock_setup_sentry.assert_called_once_with("https://user@password.mock.dsn/123456", 1.0, 1.0)
+        mock_setup_sentry.assert_called_once_with(
+            "https://user@password.mock.dsn/123456", 1.0, 1.0
+        )
 
 
-def test_client_setup_sentry_not_called_when_sentry_dsn_not_exists(mock_diode_authentication):
+def test_client_setup_sentry_not_called_when_sentry_dsn_not_exists(
+    mock_diode_authentication,
+):
     """Check that DiodeClient._setup_sentry() is not called when sentry_dsn does not exist."""
     client = DiodeClient(
         target="grpc://localhost:8081",
@@ -471,7 +484,10 @@ def continuation(x, _):
         None,
     )
     request = None
-    assert interceptor.intercept_unary_unary(continuation, client_call_details, request) == "/my/path/diode.v1.IngesterService/Ingest"
+    assert (
+        interceptor.intercept_unary_unary(continuation, client_call_details, request)
+        == "/my/path/diode.v1.IngesterService/Ingest"
+    )
 
 
 def test_interceptor_intercepts_stream_unary_calls():
@@ -491,7 +507,9 @@ def continuation(x, _):
     )
     request_iterator = None
     assert (
-        interceptor.intercept_stream_unary(continuation, client_call_details, request_iterator)
+        interceptor.intercept_stream_unary(
+            continuation, client_call_details, request_iterator
+        )
         == "/my/path/diode.v1.IngesterService/Ingest"
     )
 
@@ -569,7 +587,9 @@ def test_diode_authentication_success(mock_diode_authentication):
     with mock.patch("http.client.HTTPConnection") as mock_http_conn:
         mock_conn_instance = mock_http_conn.return_value
         mock_conn_instance.getresponse.return_value.status = 200
-        mock_conn_instance.getresponse.return_value.read.return_value = json.dumps({"access_token": "mocked_token"}).encode()
+        mock_conn_instance.getresponse.return_value.read.return_value = json.dumps(
+            {"access_token": "mocked_token"}
+        ).encode()
 
         token = auth.authenticate()
         assert token == "mocked_token"
@@ -595,14 +615,17 @@ def test_diode_authentication_failure(mock_diode_authentication):
         assert "Failed to obtain access token" in str(excinfo.value)
 
 
-@pytest.mark.parametrize("path", [
-    "/diode",
-    "",
-    None,
-    "/diode/",
-    "diode",
-    "diode/",
-    ])
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/diode",
+        "",
+        None,
+        "/diode/",
+        "diode",
+        "diode/",
+    ],
+)
 def test_diode_authentication_url_with_path(mock_diode_authentication, path):
     """Test that the authentication URL is correctly formatted with a path."""
     auth = _DiodeAuthentication(
@@ -616,9 +639,13 @@ def test_diode_authentication_url_with_path(mock_diode_authentication, path):
     with mock.patch("http.client.HTTPConnection") as mock_http_conn:
         mock_conn_instance = mock_http_conn.return_value
         mock_conn_instance.getresponse.return_value.status = 200
-        mock_conn_instance.getresponse.return_value.read.return_value = json.dumps({"access_token": "mocked_token"}).encode()
+        mock_conn_instance.getresponse.return_value.read.return_value = json.dumps(
+            {"access_token": "mocked_token"}
+        ).encode()
         auth.authenticate()
-        mock_conn_instance.request.assert_called_once_with("POST", f"{(path or '').rstrip('/')}/auth/token", mock.ANY, mock.ANY)
+        mock_conn_instance.request.assert_called_once_with(
+            "POST", f"{(path or '').rstrip('/')}/auth/token", mock.ANY, mock.ANY
+        )
 
 
 def test_diode_authentication_request_exception(mock_diode_authentication):
@@ -639,6 +666,7 @@ def test_diode_authentication_request_exception(mock_diode_authentication):
             auth.authenticate()
         assert "Failed to obtain access token: Connection error" in str(excinfo.value)
 
+
 def test_ingest_dry_run_stdout(capsys):
     """Verify ingest prints JSON when dry run is enabled."""
     client = DiodeDryRunClient()
@@ -664,6 +692,7 @@ def test_ingest_dry_run_file(tmp_path):
     assert client._stub.Ingest.call_count == 0
     assert output_file.read_text().startswith("{")
 
+
 def test_load_dryrun_entities(tmp_path):
     """Verify ``load_dryrun_entities`` yields protobuf entities."""
     output_file = tmp_path / "dryrun.jsonl"
@@ -674,7 +703,7 @@ def test_load_dryrun_entities(tmp_path):
     entities = list(load_dryrun_entities(output_file))
 
     assert len(entities) == 2
-    assert isinstance(entities[0], EntityPb)
+    assert isinstance(entities[0], ingester_pb2.Entity)
     assert entities[0].site.name == "Site1"
-    assert isinstance(entities[1], EntityPb)
+    assert isinstance(entities[1], ingester_pb2.Entity)
     assert entities[1].device.name == "Device1"