Adding unit test for safe_cupy_array function

Signed-off-by: Hrishith Thadicherla <[email protected]>

Author: hthadicherla

tests/gpu/onnx/test_quantize_onnx_torch_int4_awq.py

@@ -39,56 +39,38 @@
 #       test_qdq_utils_fp8.py::test_fused_q[bf16,fp16] fails if this script runs after the int4 test, but not before.
 
 
-def test_safe_cupy_array_all_paths(monkeypatch):
-    """Test safe_cupy_array covering all code paths including ml_dtypes handling"""
-    # Test 1: When ml_dtypes import fails (covers ImportError path)
-    # Temporarily remove ml_dtypes from sys.modules
-    import sys
-
-    if "ml_dtypes" in sys.modules:
-        ml_dtypes_backup = sys.modules["ml_dtypes"]
-        monkeypatch.delitem(sys.modules, "ml_dtypes")
-    else:
-        ml_dtypes_backup = None
-
-    tensor = np.array([1, 2, 3, 4], dtype=np.int8)
-    result = int4.safe_cupy_array(tensor)
-    assert isinstance(result, np.ndarray)  # Should return numpy array
-
-    # Restore ml_dtypes if it existed
-    if ml_dtypes_backup:
-        sys.modules["ml_dtypes"] = ml_dtypes_backup
-
-    # Test 2: When ml_dtypes exists and tensor has ml_dtypes.int4 dtype
+def test_safe_cupy_array(monkeypatch):
+    """Comprehensive test for safe_cupy_array covering all code paths."""
+    import builtins
+    import numpy  # Import actual numpy for creating int4 tensors
+    
+    # Test 1: Regular numpy array (should hit line 122)
+    result = int4.safe_cupy_array(numpy.array([1, 2, 3, 4], dtype=numpy.float32))
+    assert isinstance(result, np.ndarray)
+    
+    # Test 2: With real ml_dtypes.int4 (covers lines 117-118)
     try:
         import ml_dtypes
-
-        # Create a mock tensor with int4 dtype
-        class MockInt4Tensor:
-            def __init__(self, data):
-                self.data = data
-                self.dtype = ml_dtypes.int4
-                self.shape = data.shape
-
-            def astype(self, dtype):
-                return self.data.astype(dtype)
-
-            def __array__(self):
-                return self.data
-
-        mock_tensor = MockInt4Tensor(np.array([1, 2, 3, 4], dtype=np.int8))
-        result = int4.safe_cupy_array(mock_tensor)
-        assert isinstance(result, np.ndarray)
-        assert result.dtype == np.int8
+        int4_tensor = numpy.array([1, 2, -3, 4], dtype=numpy.float32).astype(ml_dtypes.int4)
+        result = int4.safe_cupy_array(int4_tensor)
+        assert isinstance(result, np.ndarray) and result.dtype == numpy.int8
+        expected = int4_tensor.astype(numpy.int8)
+        actual = result.get() if int4.has_cupy else result
+        np.testing.assert_array_equal(actual, expected)
     except ImportError:
-        # ml_dtypes not available, skip this part
-        pass
-
-    # Test 3: Normal case with regular numpy array
-    tensor = np.array([1, 2, 3, 4], dtype=np.int8)
-    result = int4.safe_cupy_array(tensor)
-    # Should work normally
-    assert isinstance(result, (np.ndarray, type(tensor)))
+        pass  # ml_dtypes not available
+    
+    # Test 3: When ml_dtypes import fails (covers ImportError catch and line 122)
+    def mock_import(name, *args, **kwargs):
+        if name == "ml_dtypes":
+            raise ImportError("ml_dtypes not available")
+        return builtins.__import__(name, *args, **kwargs)
+    
+    monkeypatch.setattr(builtins, "__import__", mock_import)
+    
+    # Use actual numpy for creating the array
+    result = int4.safe_cupy_array(numpy.array([5, 6, 7, 8], dtype=numpy.int8))
+    assert isinstance(result, np.ndarray)
 
 
 def test_int4_awq(tmp_path):