"""
Unit tests for the rp_cuda module
"""

from unittest.mock import patch

from runpod.serverless.utils import rp_cuda


def test_is_available_true():
    """
    Test that is_available returns True when nvidia-smi is available
    """
    with patch(
        "subprocess.check_output", return_value=b"NVIDIA-SMI"
    ) as mock_check_output:
        assert rp_cuda.is_available() is True
    mock_check_output.assert_called_once_with("nvidia-smi", shell=True)


def test_is_available_false():
    """
    Test that is_available returns False when nvidia-smi is not available
    """
    with patch(
        "subprocess.check_output", return_value=b"Not a GPU output"
    ) as mock_check_output:
        assert rp_cuda.is_available() is False
    mock_check_output.assert_called_once_with("nvidia-smi", shell=True)


def test_is_available_exception():
    """
    Test that is_available returns False when nvidia-smi raises an exception
    """
    with patch(
        "subprocess.check_output", side_effect=Exception("Bad Command")
    ) as mock_check:
        assert rp_cuda.is_available() is False
    mock_check.assert_called_once_with("nvidia-smi", shell=True)