From 1e7bfe6e24139430c816f1495cdaf8ebc78f3980 Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Thu, 7 Mar 2024 11:05:49 -0800 Subject: [PATCH] fix testing in ci --- ci/test_python.sh | 7 +- ci/test_wheel.sh | 5 - .../tests/test_cpu_handlers.py | 37 +++++++ .../tests/test_gpu_handlers.py | 80 ++++++++++++++ jupyterlab_nvdashboard/tests/test_handlers.py | 103 ------------------ 5 files changed, 118 insertions(+), 114 deletions(-) create mode 100644 jupyterlab_nvdashboard/tests/test_cpu_handlers.py create mode 100644 jupyterlab_nvdashboard/tests/test_gpu_handlers.py delete mode 100644 jupyterlab_nvdashboard/tests/test_handlers.py diff --git a/ci/test_python.sh b/ci/test_python.sh index 4eebd81..13550ed 100755 --- a/ci/test_python.sh +++ b/ci/test_python.sh @@ -18,7 +18,7 @@ set +u conda activate test set -u -# rapids-logger "Downloading artifacts from previous jobs" +rapids-logger "Downloading artifacts from previous jobs" PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python) rapids-print-env @@ -35,12 +35,7 @@ trap "EXITCODE=1" ERR set +e rapids-logger "pytest jupyterlab-nvdashboard" -# Start JupyterLab in the background -python -m jupyterlab --no-browser --port=8888 --allow-root --NotebookApp.token='' --NotebookApp.password='' & -JUPYTER_PID=$! JUPYTER_PLATFORM_DIRS=1 python -m pytest -# Shut down JupyterLab -kill $JUPYTER_PID rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh index acae426..8497bd5 100755 --- a/ci/test_wheel.sh +++ b/ci/test_wheel.sh @@ -20,12 +20,7 @@ trap "EXITCODE=1" ERR set +e rapids-logger "pytest jupyterlab-nvdashboard" -# Start JupyterLab in the background -python -m jupyterlab --no-browser --port=8888 --allow-root --NotebookApp.token='' --NotebookApp.password='' & -JUPYTER_PID=$! JUPYTER_PLATFORM_DIRS=1 python -m pytest -# Shut down JupyterLab -kill $JUPYTER_PID rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/jupyterlab_nvdashboard/tests/test_cpu_handlers.py b/jupyterlab_nvdashboard/tests/test_cpu_handlers.py new file mode 100644 index 0000000..53ce9b5 --- /dev/null +++ b/jupyterlab_nvdashboard/tests/test_cpu_handlers.py @@ -0,0 +1,37 @@ +import json +import pytest +from unittest.mock import MagicMock, patch + +from jupyterlab_nvdashboard.apps.cpu import CPUResourceWebSocketHandler + + +@pytest.fixture +def mock_handler(monkeypatch): + mock = MagicMock() + monkeypatch.setattr( + "jupyterlab_nvdashboard.apps.cpu.CustomWebSocketHandler.write_message", + mock, + ) + return mock + + +@pytest.fixture +def handler_args(): + with patch("tornado.web.Application") as mock_application, patch( + "tornado.httputil.HTTPServerRequest" + ) as mock_request: + yield mock_application, mock_request + + +def test_cpu_resource_handler(mock_handler, handler_args): + handler = CPUResourceWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "time" in data + assert "cpu_utilization" in data + assert "memory_usage" in data + assert "disk_read" in data + assert "disk_write" in data + assert "network_read" in data + assert "network_write" in data diff --git a/jupyterlab_nvdashboard/tests/test_gpu_handlers.py b/jupyterlab_nvdashboard/tests/test_gpu_handlers.py new file mode 100644 index 0000000..c204a2e --- /dev/null +++ b/jupyterlab_nvdashboard/tests/test_gpu_handlers.py @@ -0,0 +1,80 @@ +import json +import pytest +from unittest.mock import MagicMock, patch + +from jupyterlab_nvdashboard.apps.gpu import ( + GPUUtilizationWebSocketHandler, + GPUUsageWebSocketHandler, + GPUResourceWebSocketHandler, + NVLinkThroughputWebSocketHandler, + PCIStatsWebSocketHandler, +) + + +@pytest.fixture +def mock_handler(monkeypatch): + mock = MagicMock() + monkeypatch.setattr( + "jupyterlab_nvdashboard.apps.gpu.CustomWebSocketHandler.write_message", + mock, + ) + return mock + + +@pytest.fixture +def handler_args(): + with patch("tornado.web.Application") as mock_application, patch( + "tornado.httputil.HTTPServerRequest" + ) as mock_request: + yield mock_application, mock_request + + +def test_gpu_utilization_handler(mock_handler, handler_args): + handler = GPUUtilizationWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "gpu_utilization" in data + + +def test_gpu_usage_handler(mock_handler, handler_args): + handler = GPUUsageWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "memory_usage" in data + assert "total_memory" in data + + +def test_gpu_resource_handler(mock_handler, handler_args): + handler = GPUResourceWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "time" in data + assert "gpu_utilization_total" in data + assert "gpu_memory_total" in data + assert "rx_total" in data + assert "tx_total" in data + assert "gpu_memory_individual" in data + assert "gpu_utilization_individual" in data + + +def test_nvlink_throughput_handler(mock_handler, handler_args): + handler = NVLinkThroughputWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "nvlink_rx" in data + assert "nvlink_tx" in data + assert "max_rxtx_bw" in data + + +def test_pci_stats_handler(mock_handler, handler_args): + handler = PCIStatsWebSocketHandler(*handler_args) + handler.send_data() + args, _ = mock_handler.call_args + data = json.loads(args[0]) + assert "pci_tx" in data + assert "pci_rx" in data + assert "max_rxtx_tp" in data diff --git a/jupyterlab_nvdashboard/tests/test_handlers.py b/jupyterlab_nvdashboard/tests/test_handlers.py deleted file mode 100644 index 2c31533..0000000 --- a/jupyterlab_nvdashboard/tests/test_handlers.py +++ /dev/null @@ -1,103 +0,0 @@ -import json -import pytest -import websockets - -URL_PATH = "ws://localhost:8888/nvdashboard" - - -@pytest.mark.asyncio -async def test_gpu_utilization_handler(): - async with websockets.connect(f"{URL_PATH}/gpu_utilization") as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - # Receive the actual data message - response = await websocket.recv() - data = json.loads(response) - assert "gpu_utilization" in data - - -@pytest.mark.asyncio -async def test_gpu_usage_handler(): - async with websockets.connect(f"{URL_PATH}/gpu_usage") as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - response = await websocket.recv() - data = json.loads(response) - assert "memory_usage" in data - assert "total_memory" in data - - -@pytest.mark.asyncio -async def test_gpu_resource_handler(): - async with websockets.connect(f"{URL_PATH}/gpu_resource") as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - response = await websocket.recv() - data = json.loads(response) - assert "time" in data - assert "gpu_utilization_total" in data - assert "gpu_memory_total" in data - assert "rx_total" in data - assert "tx_total" in data - assert "gpu_memory_individual" in data - assert "gpu_utilization_individual" in data - - -@pytest.mark.asyncio -async def test_pci_stats_handler(): - async with websockets.connect(f"{URL_PATH}/pci_stats") as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - response = await websocket.recv() - data = json.loads(response) - assert "pci_tx" in data - assert "pci_rx" in data - assert "max_rxtx_tp" in data - - -@pytest.mark.asyncio -async def test_nvlink_throughput_handler(): - async with websockets.connect( - f"{URL_PATH}/nvlink_throughput" - ) as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - response = await websocket.recv() - data = json.loads(response) - assert "nvlink_rx" in data - assert "nvlink_tx" in data - assert "max_rxtx_bw" in data - - -@pytest.mark.asyncio -async def test_cpu_handlers(): - async with websockets.connect(f"{URL_PATH}/cpu_resource") as websocket: - # Receive the initial status message - status_message = await websocket.recv() - status_data = json.loads(status_message) - assert status_data.get("status") == "connected" - - response = await websocket.recv() - data = json.loads(response) - assert "time" in data - assert "cpu_utilization" in data - assert "memory_usage" in data - assert "disk_read" in data - assert "disk_write" in data - assert "network_read" in data - assert "network_write" in data