From 03984b4a485e8e0634dcbde7bae72bac00035981 Mon Sep 17 00:00:00 2001
From: Peter Krull - ANSYS <peter.krull@ansys.com>
Date: Fri, 22 Nov 2024 08:36:56 -0700
Subject: [PATCH] Use `weakref.finalize()` to shutdown server connection (#580)

---
 pyproject.toml                                |  2 +-
 .../server_connection/server_connection.py    | 13 +++-
 .../test_server_connection.py                 | 66 +++++++++++++++++++
 3 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 15f47615a..d09434b88 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "flit_core.buildapi"
 [project]
 # Check https://flit.readthedocs.io/en/latest/pyproject_toml.html for all available sections
 name = "ansys-additive-core"
-version = "0.20.dev1"
+version = "0.20.dev2"
 description = "A Python client for the Ansys Additive service"
 readme = "README.rst"
 requires-python = ">=3.10,<4"
diff --git a/src/ansys/additive/core/server_connection/server_connection.py b/src/ansys/additive/core/server_connection/server_connection.py
index e0499da9f..353a86e56 100644
--- a/src/ansys/additive/core/server_connection/server_connection.py
+++ b/src/ansys/additive/core/server_connection/server_connection.py
@@ -24,6 +24,7 @@
 import logging
 import os
 import time
+import weakref
 from dataclasses import dataclass
 
 import grpc
@@ -118,6 +119,8 @@ def __init__(
     ) -> None:
         """Initialize a server connection."""
 
+        weakref.finalize(self, self.disconnect)
+
         if channel and addr:
             raise ValueError("Both 'channel' and 'addr' cannot both be specified.")
 
@@ -158,12 +161,18 @@ def __init__(
 
         self._log.info("Connected to %s", self.channel_str)
 
-    def __del__(self):
-        """Destructor for cleaning up server connection."""
+    def __del__(self) -> None:
+        """Disconnect from server."""
+        self.disconnect()
+
+    def disconnect(self):
+        """Clean up server connection."""
         if hasattr(self, "_server_instance") and self._server_instance:
             self._server_instance.delete()
+            self._server_instance = None
         if hasattr(self, "_server_process") and self._server_process:
             self._server_process.kill()
+            self._server_process = None
 
     @property
     def channel_str(self) -> str:
diff --git a/tests/server_connection/test_server_connection.py b/tests/server_connection/test_server_connection.py
index 667d42cb1..856470196 100644
--- a/tests/server_connection/test_server_connection.py
+++ b/tests/server_connection/test_server_connection.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+from subprocess import Popen
 from unittest.mock import ANY, Mock, create_autospec
 
 import grpc
@@ -347,3 +348,68 @@ def test_server_connection_status_str_connected_with_metadata():
     # assert
     expected_str = "Server localhost:50051 is connected.\n  version: 1.0\n  status: running"
     assert status_str == expected_str
+
+
+def test_disconnect_with_server_instance(monkeypatch):
+    # arrange
+    mock_ready = create_autospec(
+        ansys.additive.core.server_connection.server_connection.ServerConnection.ready,
+        return_value=True,
+    )
+    monkeypatch.setattr(
+        ansys.additive.core.server_connection.server_connection.ServerConnection,
+        "ready",
+        mock_ready,
+    )
+    mock_server_instance = Mock(pypim.Instance)
+    server = ServerConnection(channel=grpc.insecure_channel("target"))
+    server._server_instance = mock_server_instance
+
+    # act
+    server.disconnect()
+
+    # assert
+    mock_server_instance.delete.assert_called_once()
+
+
+def test_disconnect_with_server_process(monkeypatch):
+    # arrange
+    mock_ready = create_autospec(
+        ansys.additive.core.server_connection.server_connection.ServerConnection.ready,
+        return_value=True,
+    )
+    monkeypatch.setattr(
+        ansys.additive.core.server_connection.server_connection.ServerConnection,
+        "ready",
+        mock_ready,
+    )
+    mock_server_process = Mock(Popen)
+    server = ServerConnection(channel=grpc.insecure_channel("target"))
+    server._server_process = mock_server_process
+
+    # act
+    server.disconnect()
+
+    # assert
+    mock_server_process.kill.assert_called_once()
+
+
+def test_disconnect_with_no_server_instance_or_process(monkeypatch):
+    # arrange
+    mock_ready = create_autospec(
+        ansys.additive.core.server_connection.server_connection.ServerConnection.ready,
+        return_value=True,
+    )
+    monkeypatch.setattr(
+        ansys.additive.core.server_connection.server_connection.ServerConnection,
+        "ready",
+        mock_ready,
+    )
+    server = ServerConnection(channel=grpc.insecure_channel("target"))
+
+    # act
+    server.disconnect()
+
+    # assert
+    # No exception should be raised and no methods should be called
+    assert True