Merge pull request #184 from kbakk/178-grpc-conf-port

Configure gRPC channel (port from 2.3.2 to 3.0.0)

pyzeebe/grpc_internals/channel_options.py

@@ -0,0 +1,37 @@
+""" gRPC Channel Options
+
+``grpc.keepalive_time_ms``
+--------------------------
+
+Time between keepalive pings. Following the official Zeebe Java/Go client, sending pings every 45 seconds.
+
+https://docs.camunda.io/docs/product-manuals/zeebe/deployment-guide/operations/setting-up-a-cluster/#keep-alive-intervals
+"""
+from typing import Any, Dict, Tuple
+
+GRPC_CHANNEL_OPTIONS = {
+    "grpc.keepalive_time_ms": 45_000
+}
+
+
+def get_channel_options(options: Dict[str, Any] = None) -> Tuple[Tuple[str, Any], ...]:
+    """
+    Convert options dict to tuple in expected format for creating the gRPC channel
+
+    Args:
+        options (Dict[str, Any]): A key/value representation of `gRPC channel arguments_`.
+            Default: None (will use library defaults)
+
+    Returns:
+        Tuple[Tuple[str, Any], ...]: Options for the gRPC channel
+
+    .. _gRPC channel arguments:
+        https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
+    """
+    if options:
+        options = {**GRPC_CHANNEL_OPTIONS, **options}
+    else:
+        options = GRPC_CHANNEL_OPTIONS
+    return tuple(
+        (k, v) for k, v in options.items()
+    )

pyzeebe/grpc_internals/grpc_channel_utils.py

@@ -1,9 +1,10 @@
 import os
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import grpc
 
 from pyzeebe.credentials.base_credentials import BaseCredentials
+from pyzeebe.grpc_internals.channel_options import get_channel_options
 
 
 def create_connection_uri(
@@ -21,10 +22,18 @@ def create_connection_uri(
 def create_channel(
     connection_uri: str,
     credentials: Optional[BaseCredentials] = None,
-    secure_connection: bool = False
+    secure_connection: bool = False,
+    options: Dict[str, Any] = None
 ) -> grpc.aio.Channel:
+    """
+    options: A key/value representation of `gRPC channel arguments_`. Default: None (will use library defaults)
+
+    .. _gRPC channel arguments:
+        https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
+    """
+    channel_options = get_channel_options(options)
     if credentials:
-        return grpc.aio.secure_channel(connection_uri, credentials.grpc_credentials)
+        return grpc.aio.secure_channel(connection_uri, credentials.grpc_credentials, options=channel_options)
     if secure_connection:
-        return grpc.aio.secure_channel(connection_uri, grpc.ssl_channel_credentials())
-    return grpc.aio.insecure_channel(connection_uri)
+        return grpc.aio.secure_channel(connection_uri, grpc.ssl_channel_credentials(), options=channel_options)
+    return grpc.aio.insecure_channel(connection_uri, options=channel_options)

tests/unit/grpc_internals/channel_options_test.py

@@ -0,0 +1,69 @@
+from copy import deepcopy
+from unittest.mock import Mock, patch
+
+import pytest
+
+import pyzeebe.grpc_internals.channel_options
+from pyzeebe.grpc_internals.channel_options import get_channel_options
+from pyzeebe.grpc_internals.zeebe_adapter_base import ZeebeAdapterBase
+
+
+@pytest.fixture
+def revert_monkeypatch_after_test():
+    """
+    This sort of exists in pytest already (docs.pytest.org/en/stable/monkeypatch.html),
+    however that means a bit of "magic" happens, this is a bit clearer and tests the users
+    approach to this.
+    """
+    options_before = deepcopy(pyzeebe.grpc_internals.channel_options.GRPC_CHANNEL_OPTIONS)
+    yield
+    pyzeebe.grpc_internals.channel_options.GRPC_CHANNEL_OPTIONS = options_before
+
+
+def test_get_channel_options_returns_tuple_of_tuple_with_options():
+    assert get_channel_options() == (
+        ("grpc.keepalive_time_ms", 45000),
+    )
+
+
+@pytest.mark.parametrize("grpc_method,call_kwargs",
+                         [
+                             ("grpc.aio.secure_channel", {"secure_connection": True}),
+                             ("grpc.aio.insecure_channel", {}),
+                             ("grpc.aio.secure_channel", {"credentials": Mock()}),
+                         ])
+def test_create_channel_called_with_options(grpc_method, call_kwargs, zeebe_adapter):
+    with patch(grpc_method) as channel_mock:
+        ZeebeAdapterBase(**call_kwargs).connect()
+        expected_options = (('grpc.keepalive_time_ms', 45000),)
+        # `call_args.kwargs` as it's not available in python <=3.7
+        # https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.call_args
+        # 0 is args, 1 is kwargs
+        assert channel_mock.call_args[1]["options"] == expected_options
+
+
+@pytest.mark.usefixtures("revert_monkeypatch_after_test")
+def test_monkeypatching_with_options_override():
+    pyzeebe.grpc_internals.channel_options.GRPC_CHANNEL_OPTIONS["grpc.keepalive_time_ms"] = 4000
+    assert get_channel_options() == (
+        ("grpc.keepalive_time_ms", 4000),
+    )
+
+
+@pytest.mark.usefixtures("revert_monkeypatch_after_test")
+def test_monkeypatching_with_options_added():
+    pyzeebe.grpc_internals.channel_options.GRPC_CHANNEL_OPTIONS.update({
+        "grpc.keepalive_timeout_ms": 120000,
+        "grpc.http2.min_time_between_pings_ms": 60000,
+    })
+    assert get_channel_options() == (
+        ("grpc.keepalive_time_ms", 45000),
+        ("grpc.keepalive_timeout_ms", 120000),
+        ("grpc.http2.min_time_between_pings_ms", 60000)
+    )
+
+
+@pytest.mark.usefixtures("revert_monkeypatch_after_test")
+def test_monkeypatching_with_options_removed():
+    pyzeebe.grpc_internals.channel_options.GRPC_CHANNEL_OPTIONS = {}
+    assert get_channel_options() == ()

tests/unit/grpc_internals/grpc_channel_utils_test.py

@@ -26,6 +26,29 @@ def test_creates_secure_channel_when_secure_connection_is_enabled(self):
 
             channel_mock.assert_called_once()
 
+    def test_creates_secure_channel_with_default_options(self):
+        with patch("grpc.aio.insecure_channel") as channel_mock:
+            create_channel(str(uuid4()), options=None)
+
+            expected_options = (
+                ("grpc.keepalive_time_ms", 45000),
+            )
+            assert channel_mock.call_args[1]["options"] == expected_options
+
+    def test_creates_insecure_channel_with_options(self):
+        with patch("grpc.aio.insecure_channel") as channel_mock:
+            additional_options = {
+                "grpc.keepalive_timeout_ms": 120000,
+                "grpc.http2.min_time_between_pings_ms": 60000,
+            }
+            create_channel(str(uuid4()), options=additional_options)
+
+            expected_options = (
+                ("grpc.keepalive_time_ms", 45000),
+                ("grpc.keepalive_timeout_ms", 120000),
+                ("grpc.http2.min_time_between_pings_ms", 60000),
+            )
+            assert channel_mock.call_args[1]["options"] == expected_options
 
 class TestCreateConnectionUri:
     def test_uses_credentials_first(self):