add tests.

dmlc · Jul 21, 2024 · 060001d · 060001d
1 parent cad4b6e
commit 060001d
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests/python/pytorch/graphbolt/impl/test_basic_feature_store.py b/tests/python/pytorch/graphbolt/impl/test_basic_feature_store.py
@@ -15,6 +15,14 @@ def test_basic_feature_store_homo():
 
     feature_store = gb.BasicFeatureStore(features)
 
+    # Test __getitem__ to access the stored Feature.
+    feature = feature_store[("node", None, "a")]
+    assert isinstance(feature, gb.Feature)
+    assert torch.equal(
+        feature.read(),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+
     # Test read the entire feature.
     assert torch.equal(
         feature_store.read("node", None, "a"),
@@ -60,6 +68,14 @@ def test_basic_feature_store_hetero():
 
     feature_store = gb.BasicFeatureStore(features)
 
+    # Test __getitem__ to access the stored Feature.
+    feature = feature_store[("node", "author", "a")]
+    assert isinstance(feature, gb.Feature)
+    assert torch.equal(
+        feature.read(),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+
     # Test read the entire feature.
     assert torch.equal(
         feature_store.read("node", "author", "a"),