From 60bdf8514db11166821e2a1b5f97306fa3427871 Mon Sep 17 00:00:00 2001
From: Deezzir <deezzir@gmail.com>
Date: Fri, 13 Sep 2024 13:46:52 -0400
Subject: [PATCH] Refinements

---
 snap/hooks/configure               |  2 ++
 snap/local/run_dcgm_exporter.sh    |  2 +-
 tests/functional/conftest.py       | 17 +++++++++++++++++
 tests/functional/requirments.txt   |  4 ++--
 tests/functional/test_snap_dcgm.py | 17 +++--------------
 5 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/snap/hooks/configure b/snap/hooks/configure
index f05078a..a55fc8a 100644
--- a/snap/hooks/configure
+++ b/snap/hooks/configure
@@ -4,9 +4,11 @@
 # so users can see available options by running
 # `sudo snap get dcgm`.
 if [ -z "$(snapctl get nv-hostengine-port)" ]; then
+    # Setting to empty string for the nv-hostengine binary to use the default port. (5555)
     snapctl set nv-hostengine-port=""
 fi
 
 if [ -z "$(snapctl get dcgm-exporter-address)" ]; then
+    # Setting to empty string for the dcgm-exporter binary to use the default address. (:9400)
     snapctl set dcgm-exporter-address=""
 fi
diff --git a/snap/local/run_dcgm_exporter.sh b/snap/local/run_dcgm_exporter.sh
index e852dd9..a1dd641 100644
--- a/snap/local/run_dcgm_exporter.sh
+++ b/snap/local/run_dcgm_exporter.sh
@@ -4,7 +4,7 @@ set -euo pipefail
 # Build the argument list for the dcgm-exporter command
 args=()
 
-# Add the dcgm-exporter-port option if it is set. Default: “:9400”
+# Add the dcgm-exporter-address option if it is set. Default: “:9400”
 dcgm_exporter_address="$(snapctl get dcgm-exporter-address)"
 
 if [ -n "$dcgm_exporter_address" ]; then
diff --git a/tests/functional/conftest.py b/tests/functional/conftest.py
index 56b80a0..5e3e58a 100644
--- a/tests/functional/conftest.py
+++ b/tests/functional/conftest.py
@@ -1,6 +1,20 @@
 import subprocess
 
 import pytest
+from tenacity import retry, stop_after_delay, wait_fixed
+
+
+@retry(wait=wait_fixed(5), stop=stop_after_delay(20))
+def check_dcgm_exporter_service():
+    dcgm_exporter_service = "dcgm.dcgm-exporter"
+
+    result = subprocess.run(
+        f"sudo snap services {dcgm_exporter_service}".split(),
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    assert " active" in result.stdout.strip(), f"{dcgm_exporter_service} service is not active"
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -15,6 +29,9 @@ def install_dcgm_snap():
         shell=True,
     )
 
+    subprocess.run("sudo snap start dcgm.dcgm-exporter".split(), check=True)
+    check_dcgm_exporter_service()
+
     yield
 
     subprocess.run("sudo snap remove --purge dcgm".split(), check=True)
diff --git a/tests/functional/requirments.txt b/tests/functional/requirments.txt
index 771e68b..bd2d4ef 100644
--- a/tests/functional/requirments.txt
+++ b/tests/functional/requirments.txt
@@ -1,2 +1,2 @@
-tenacity==9.0.0
-pyyaml==6.0.2
+tenacity
+pyyaml
diff --git a/tests/functional/test_snap_dcgm.py b/tests/functional/test_snap_dcgm.py
index 842d6df..7d38f56 100644
--- a/tests/functional/test_snap_dcgm.py
+++ b/tests/functional/test_snap_dcgm.py
@@ -6,21 +6,10 @@
 from tenacity import retry, stop_after_delay, wait_fixed
 
 
-@retry(wait=wait_fixed(5), stop=stop_after_delay(15))
+@retry(wait=wait_fixed(5), stop=stop_after_delay(30))
 def test_dcgm_exporter():
     """Test of the dcgm-exporter service and its endpoint."""
     endpoint = "http://localhost:9400/metrics"
-    dcgm_exporter_service = "dcgm.dcgm-exporter"
-
-    subprocess.run(f"sudo snap start {dcgm_exporter_service}".split(), check=True)
-
-    result = subprocess.run(
-        f"sudo snap services {dcgm_exporter_service}".split(),
-        check=True,
-        capture_output=True,
-        text=True,
-    )
-    assert " active" in result.stdout.strip(), "dcgm-exporter service is not active"
 
     # Check the exporter endpoint, will raise an exception if the endpoint is not reachable
     urllib.request.urlopen(endpoint)
@@ -37,7 +26,7 @@ def test_dcgm_nv_hostengine():
         text=True,
     )
 
-    assert " active" in service.stdout.strip(), "nv-hostengine service is not active"
+    assert " active" in service.stdout.strip(), f"{nv_hostengine_service} service is not active"
 
 
 def test_dcgmi():
@@ -48,7 +37,7 @@ def test_dcgmi():
     assert "GPU ID" in result.stdout.strip(), "DCGMI is not working"
 
 
-def test_dcgm_port_configs():
+def test_dcgm_bind_configs():
     """Test snap port configuratin."""
     services = ["dcgm.dcgm-exporter", "dcgm.nv-hostengine"]
     configs = ["dcgm-exporter-address", "nv-hostengine-port"]