refactor to use PanelClient

Author: Mike Prosser

src/nipanel/_panel.py

@@ -1,15 +1,13 @@
 from __future__ import annotations
 
 import sys
-import time
 from abc import ABC, abstractmethod
 from types import TracebackType
 from typing import TYPE_CHECKING, Optional, Type
 
-from grpc import RpcError
-from ni.pythonpanel.v1.python_panel_service_pb2 import ConnectRequest, DisconnectRequest
-from ni.pythonpanel.v1.python_panel_service_pb2_grpc import PythonPanelServiceStub
-from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
+from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
+
+from nipanel._panel_client import PanelClient
 
 if TYPE_CHECKING:
     if sys.version_info >= (3, 11):
@@ -21,19 +19,15 @@
 class Panel(ABC):
     """This class allows you to connect to a panel and specify values for its controls."""
 
-    RETRY_WAIT_TIME = 1  # time in seconds to wait before retrying connection
-
-    _channel_pool: GrpcChannelPool
-    _stub: PythonPanelServiceStub | None
+    _panel_client: PanelClient
     _panel_id: str
     _panel_uri: str
 
-    __slots__ = ["_channel_pool", "_stub", "_panel_id", "_panel_uri", "__weakref__"]
+    __slots__ = ["_panel_client", "_panel_id", "_panel_uri", "__weakref__"]
 
     def __init__(self, panel_id: str, panel_uri: str) -> None:
         """Initialize the panel."""
-        self._channel_pool = GrpcChannelPool()
-        self._stub = None
+        self._panel_client = PanelClient(self._resolve_service_address)
         self._panel_id = panel_id
         self._panel_uri = panel_uri
 
@@ -64,27 +58,11 @@ def __exit__(
 
     def connect(self) -> None:
         """Connect to the panel and open it."""
-        address = self._resolve_service_address()
-        channel = self._channel_pool.get_channel(address)
-        self._stub = PythonPanelServiceStub(channel)
-
-        connect_request = ConnectRequest(panel_id=self._panel_id, panel_uri=self._panel_uri)
-        try:
-            self._stub.Connect(connect_request)
-        except RpcError:
-            # retry the connection if it fails, but only once
-            time.sleep(self.RETRY_WAIT_TIME)
-            self._stub.Connect(connect_request)
+        self._panel_client.connect(self._panel_id, self._panel_uri)
 
     def disconnect(self) -> None:
         """Disconnect from the panel (does not close the panel)."""
-        if self._stub is None:
-            raise RuntimeError("connect() must be called before disconnect()")
-
-        disconnect_request = DisconnectRequest(panel_id=self._panel_id)
-        self._stub.Disconnect(disconnect_request)
-        self._stub = None
-        self._channel_pool.close()
+        self._panel_client.disconnect(self._panel_id)
 
     def get_value(self, value_id: str) -> object:
         """Get the value for a control on the panel.
@@ -95,7 +73,7 @@ def get_value(self, value_id: str) -> object:
         Returns:
             The value
         """
-        # TODO: AB#3095681 - get the Any from _stub.GetValue and convert it to the correct type
+        # TODO: AB#3095681 - get the Any from _client.get_value and convert it to the correct type
         return "placeholder value"
 
     def set_value(self, value_id: str, value: object) -> None:
@@ -105,10 +83,10 @@ def set_value(self, value_id: str, value: object) -> None:
             value_id: The id of the value
             value: The value
         """
-        # TODO: AB#3095681 - Convert the value to an Any and pass it to _stub.SetValue
+        # TODO: AB#3095681 - Convert the value to an Any and pass it to _client.set_value
         pass
 
     @abstractmethod
-    def _resolve_service_address(self) -> str:
-        """Resolve the service location for the panel."""
+    def _resolve_service_address(self, discovery_client: DiscoveryClient) -> str:
+        """Resolve the service address for the panel."""
         raise NotImplementedError

src/nipanel/_panel_client.py

@@ -0,0 +1,81 @@
+"""Client for accessing the NI Python Panel Service."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Callable
+
+import grpc
+from ni.pythonpanel.v1.python_panel_service_pb2 import ConnectRequest, DisconnectRequest
+from ni.pythonpanel.v1.python_panel_service_pb2_grpc import PythonPanelServiceStub
+from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
+from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
+
+_logger = logging.getLogger(__name__)
+
+
+class PanelClient:
+    """Client for accessing the NI Python Panel Service."""
+
+    def __init__(
+        self,
+        resolve_service_address_fn: Callable[[DiscoveryClient], str],
+        *,
+        discovery_client: DiscoveryClient | None = None,
+        grpc_channel: grpc.Channel | None = None,
+        grpc_channel_pool: GrpcChannelPool | None = None,
+    ) -> None:
+        """Initialize the panel client.
+
+        Args:
+            resolve_service_address_fn: A function to resolve the service location.
+            discovery_client: An optional discovery client.
+            grpc_channel: An optional panel gRPC channel.
+            grpc_channel_pool: An optional gRPC channel pool.
+        """
+        self._initialization_lock = threading.Lock()
+        self._resolve_service_address_fn = resolve_service_address_fn
+        self._discovery_client = discovery_client
+        self._grpc_channel_pool = grpc_channel_pool
+        self._stub: PythonPanelServiceStub | None = None
+
+        if grpc_channel is not None:
+            self._stub = PythonPanelServiceStub(grpc_channel)
+
+    def _get_stub(self) -> PythonPanelServiceStub:
+        if self._stub is None:
+            with self._initialization_lock:
+                if self._grpc_channel_pool is None:
+                    _logger.debug("Creating unshared GrpcChannelPool.")
+                    self._grpc_channel_pool = GrpcChannelPool()
+                if self._discovery_client is None:
+                    _logger.debug("Creating unshared DiscoveryClient.")
+                    self._discovery_client = DiscoveryClient(
+                        grpc_channel_pool=self._grpc_channel_pool
+                    )
+                if self._stub is None:
+                    service_address = self._resolve_service_address_fn(self._discovery_client)
+                    channel = self._grpc_channel_pool.get_channel(service_address)
+                    self._stub = PythonPanelServiceStub(channel)
+        return self._stub
+
+    def _invoke_with_retry(self, method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        """Invoke a gRPC method with retry logic."""
+        try:
+            return method(*args, **kwargs)
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.UNAVAILABLE or e.code() == grpc.StatusCode.UNKNOWN:
+                # if the service is unavailable, we can retry the connection
+                self._stub = None
+                return method(*args, **kwargs)
+
+    def connect(self, panel_id: str, panel_uri: str) -> None:
+        """Connect to the panel and open it."""
+        connect_request = ConnectRequest(panel_id=panel_id, panel_uri=panel_uri)
+        self._invoke_with_retry(self._get_stub().Connect, connect_request)
+
+    def disconnect(self, panel_id: str) -> None:
+        """Disconnect from the panel (does not close the panel)."""
+        disconnect_request = DisconnectRequest(panel_id=panel_id)
+        self._invoke_with_retry(self._get_stub().Disconnect, disconnect_request)

src/nipanel/_streamlit_panel.py

@@ -1,5 +1,4 @@
 from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
-from ni_measurement_plugin_sdk_service.grpc.channelpool import GrpcChannelPool
 
 from nipanel._panel import Panel
 
@@ -23,8 +22,8 @@ def __init__(self, panel_id: str, streamlit_script_uri: str) -> None:
         """
         super().__init__(panel_id, streamlit_script_uri)
 
-    def _resolve_service_address(self) -> str:
-        with GrpcChannelPool() as grpc_channel_pool:
-            discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
-            service_location = discovery_client.resolve_service(self.PYTHON_PANEL_SERVICE)
-            return service_location.insecure_address
+    def _resolve_service_address(self, discovery_client: DiscoveryClient) -> str:
+        service_location = discovery_client.resolve_service(
+            provided_interface=self.PYTHON_PANEL_SERVICE, service_class=self.PYTHON_PANEL_SERVICE
+        )
+        return service_location.insecure_address

tests/unit/test_panel.py

@@ -1,5 +1,3 @@
-import pytest
-
 from tests.utils._fake_python_panel_servicer import FakePythonPanelServicer
 from tests.utils._port_panel import PortPanel
 
@@ -35,16 +33,6 @@ def test___with_panel___set_value___gets_same_value(
         assert panel.get_value("test_id") == "placeholder value"
 
 
-def test___new_panel___disconnect___raises_runtime_error(
-    fake_python_panel_service: tuple[FakePythonPanelServicer, int],
-) -> None:
-    _, port = fake_python_panel_service
-    panel = PortPanel(port, "my_panel", "path/to/script")
-
-    with pytest.raises(RuntimeError):
-        panel.disconnect()
-
-
 def test___first_connect_fails___connect___gets_value(
     fake_python_panel_service: tuple[FakePythonPanelServicer, int],
 ) -> None:

tests/utils/_fake_python_panel_servicer.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import google.protobuf.any_pb2 as any_pb2
+import grpc
 from ni.pythonpanel.v1.python_panel_service_pb2 import (
     ConnectRequest,
     ConnectResponse,
@@ -24,7 +25,7 @@ def Connect(self, request: ConnectRequest, context: Any) -> ConnectResponse:  #
         """Trivial implementation for testing."""
         if self._fail_next_connect:
             self._fail_next_connect = False
-            raise ValueError("Simulate a failure to Connect.")
+            context.abort(grpc.StatusCode.UNAVAILABLE, "Simulated connection failure")
         return ConnectResponse()
 
     def Disconnect(  # noqa: N802

tests/utils/_port_panel.py

@@ -1,3 +1,5 @@
+from ni_measurement_plugin_sdk_service.discovery import DiscoveryClient
+
 from nipanel._panel import Panel
 
 
@@ -18,5 +20,5 @@ def __init__(self, port: int, panel_id: str, panel_uri: str) -> None:
         super().__init__(panel_id, panel_uri)
         self.port = port
 
-    def _resolve_service_address(self) -> str:
+    def _resolve_service_address(self, discovery_client: DiscoveryClient) -> str:
         return f"localhost:{self.port}"