add a test for Theta.pop()

nod-ai · Aug 29, 2024 · 07e6c5e · 07e6c5e
1 parent 7f3c963
commit 07e6c5e
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/sharktank/tests/types/dataset_test.py b/sharktank/tests/types/dataset_test.py
@@ -77,6 +77,22 @@ def testTransform(self):
             self.assertIsNot(pt1, pt2)
             torch.testing.assert_close(pt1, pt2)
 
+    def testPop(self):
+        t1 = Theta(
+            _flat_t_dict(
+                _t("a.b.c", 1, 2),
+                _t("a.c.d", 10, 11),
+                _t("a.b.3", 3, 4),
+            )
+        )
+        popped = t1.pop("a.b").flatten()
+        t1 = t1.flatten()
+
+        self.assertIsNotNone("a.c.d", t1.keys())
+        self.assertNotIn("a.b.c", t1.keys())
+        self.assertNotIn("a.b.3", t1.keys())
+        self.assertIn("a.b.3", popped.keys())
+
 
 class DatasetTest(unittest.TestCase):
     def setUp(self):