From 1ac0607960b78f8336f9f3d1a891258920832e66 Mon Sep 17 00:00:00 2001
From: Skeleton003 <799284168@qq.com>
Date: Fri, 2 Feb 2024 13:34:57 +0000
Subject: [PATCH] 1

---
 tests/cugraph/cugraph-ops/test_cugraph_gatconv.py   |  7 ++++---
 .../cugraph-ops/test_cugraph_relgraphconv.py        | 13 +++++++------
 tests/cugraph/cugraph-ops/test_cugraph_sageconv.py  |  7 ++++---
 tests/cugraph/test_basics.py                        |  8 ++++----
 4 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/tests/cugraph/cugraph-ops/test_cugraph_gatconv.py b/tests/cugraph/cugraph-ops/test_cugraph_gatconv.py
index a1b5524638ee..dfcafc4f87ae 100644
--- a/tests/cugraph/cugraph-ops/test_cugraph_gatconv.py
+++ b/tests/cugraph/cugraph-ops/test_cugraph_gatconv.py
@@ -2,9 +2,10 @@
 from collections import OrderedDict
 from itertools import product
 
-import dgl
 import pytest
 import torch
+
+from dgl import graph, to_block as dgl_to_block
 from dgl.nn import CuGraphGATConv, GATConv
 
 options = OrderedDict(
@@ -20,7 +21,7 @@
 def generate_graph():
     u = torch.tensor([0, 1, 0, 2, 3, 0, 4, 0, 5, 0, 6, 7, 0, 8, 9])
     v = torch.tensor([1, 9, 2, 9, 9, 4, 9, 5, 9, 6, 9, 9, 8, 9, 0])
-    g = dgl.graph((u, v))
+    g = graph((u, v))
     return g
 
 
@@ -34,7 +35,7 @@ def test_gatconv_equality(idtype_int, max_in_degree, num_heads, to_block):
     if idtype_int:
         g = g.int()
     if to_block:
-        g = dgl.to_block(g)
+        g = dgl_to_block(g)
     feat = torch.rand(g.num_src_nodes(), in_feat).to(device)
 
     torch.manual_seed(0)
diff --git a/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py b/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py
index 2dd1b9260228..050b7d37bca6 100644
--- a/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py
+++ b/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py
@@ -2,9 +2,10 @@
 from collections import OrderedDict
 from itertools import product
 
-import dgl
 import pytest
 import torch
+
+from dgl import ETYPE, graph, to_block as dgl_to_block
 from dgl.nn import CuGraphRelGraphConv, RelGraphConv
 
 # TODO(tingyu66): Re-enable the following tests after updating cuGraph CI image.
@@ -23,7 +24,7 @@
 def generate_graph():
     u = torch.tensor([0, 1, 0, 2, 3, 0, 4, 0, 5, 0, 6, 7, 0, 8, 9])
     v = torch.tensor([1, 9, 2, 9, 9, 4, 9, 5, 9, 6, 9, 9, 8, 9, 0])
-    g = dgl.graph((u, v))
+    g = graph((u, v))
     return g
 
 
@@ -41,11 +42,11 @@ def test_relgraphconv_equality(
         "self_loop": self_loop,
     }
     g = generate_graph().to(device)
-    g.edata[dgl.ETYPE] = torch.randint(num_rels, (g.num_edges(),)).to(device)
+    g.edata[ETYPE] = torch.randint(num_rels, (g.num_edges(),)).to(device)
     if idtype_int:
         g = g.int()
     if to_block:
-        g = dgl.to_block(g)
+        g = dgl_to_block(g)
     feat = torch.rand(g.num_src_nodes(), in_feat).to(device)
 
     torch.manual_seed(0)
@@ -55,8 +56,8 @@ def test_relgraphconv_equality(
     kwargs["apply_norm"] = False
     conv2 = CuGraphRelGraphConv(*args, **kwargs).to(device)
 
-    out1 = conv1(g, feat, g.edata[dgl.ETYPE])
-    out2 = conv2(g, feat, g.edata[dgl.ETYPE], max_in_degree=max_in_degree)
+    out1 = conv1(g, feat, g.edata[ETYPE])
+    out2 = conv2(g, feat, g.edata[ETYPE], max_in_degree=max_in_degree)
     assert torch.allclose(out1, out2, atol=1e-06)
 
     grad_out = torch.rand_like(out1)
diff --git a/tests/cugraph/cugraph-ops/test_cugraph_sageconv.py b/tests/cugraph/cugraph-ops/test_cugraph_sageconv.py
index cec5d3fd68ef..0aade8c29ff9 100644
--- a/tests/cugraph/cugraph-ops/test_cugraph_sageconv.py
+++ b/tests/cugraph/cugraph-ops/test_cugraph_sageconv.py
@@ -2,9 +2,10 @@
 from collections import OrderedDict
 from itertools import product
 
-import dgl
 import pytest
 import torch
+
+from dgl import graph, to_block as dgl_to_block
 from dgl.nn import CuGraphSAGEConv, SAGEConv
 
 options = OrderedDict(
@@ -19,7 +20,7 @@
 def generate_graph():
     u = torch.tensor([0, 1, 0, 2, 3, 0, 4, 0, 5, 0, 6, 7, 0, 8, 9])
     v = torch.tensor([1, 9, 2, 9, 9, 4, 9, 5, 9, 6, 9, 9, 8, 9, 0])
-    g = dgl.graph((u, v))
+    g = graph((u, v))
     return g
 
 
@@ -32,7 +33,7 @@ def test_SAGEConv_equality(idtype_int, max_in_degree, to_block):
     if idtype_int:
         g = g.int()
     if to_block:
-        g = dgl.to_block(g)
+        g = dgl_to_block(g)
     feat = torch.rand(g.num_src_nodes(), in_feat).to(device)
 
     torch.manual_seed(0)
diff --git a/tests/cugraph/test_basics.py b/tests/cugraph/test_basics.py
index 7e5b6e4ade72..2551dbd16573 100644
--- a/tests/cugraph/test_basics.py
+++ b/tests/cugraph/test_basics.py
@@ -3,7 +3,7 @@
 import cugraph  # usort: skip
 import backend as F
 
-import dgl
+from dgl import from_cugraph, graph
 
 
 def test_dummy():
@@ -12,7 +12,7 @@ def test_dummy():
 
 
 def test_to_cugraph_conversion():
-    g = dgl.graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to("cuda")
+    g = graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to("cuda")
     cugraph_g = g.to_cugraph()
 
     assert cugraph_g.number_of_nodes() == g.num_nodes()
@@ -33,7 +33,7 @@ def test_from_cugraph_conversion():
 
     cugraph_g.from_cudf_edgelist(df)
 
-    g = dgl.from_cugraph(cugraph_g)
+    g = from_cugraph(cugraph_g)
 
     assert g.device.type == "cuda"
     assert g.num_nodes() == cugraph_g.number_of_nodes()
@@ -52,7 +52,7 @@ def test_from_cugraph_conversion():
 
     cugraph_g.from_cudf_edgelist(df)
 
-    g = dgl.from_cugraph(cugraph_g)
+    g = from_cugraph(cugraph_g)
 
     assert g.device.type == "cuda"
     assert g.num_nodes() == cugraph_g.number_of_nodes()