1 files changed, 24 insertions, 1 deletions
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index 25329e2..8755b42 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -229,13 +229,36 @@ class TestLineage(unittest.TestCase):
             "output",
             "SELECT (SELECT max(t3.my_column) my_column FROM foo t3) AS output FROM table3",
         )
-        self.assertEqual(node.name, "SUBQUERY")
+        self.assertEqual(node.name, "output")
         node = node.downstream[0]
         self.assertEqual(node.name, "my_column")
         node = node.downstream[0]
         self.assertEqual(node.name, "t3.my_column")
         self.assertEqual(node.source.sql(), "foo AS t3")
 
+        node = lineage(
+            "y",
+            "SELECT SUM((SELECT max(a) a from x) + (SELECT min(b) b from x) + c) AS y FROM x",
+        )
+        self.assertEqual(node.name, "y")
+        self.assertEqual(len(node.downstream), 3)
+        self.assertEqual(node.downstream[0].name, "a")
+        self.assertEqual(node.downstream[1].name, "b")
+        self.assertEqual(node.downstream[2].name, "x.c")
+
+        node = lineage(
+            "x",
+            "WITH cte AS (SELECT a, b FROM z) SELECT sum(SELECT a FROM cte) AS x, (SELECT b FROM cte) as y FROM cte",
+        )
+        self.assertEqual(node.name, "x")
+        self.assertEqual(len(node.downstream), 1)
+        node = node.downstream[0]
+        self.assertEqual(node.name, "a")
+        node = node.downstream[0]
+        self.assertEqual(node.name, "cte.a")
+        node = node.downstream[0]
+        self.assertEqual(node.name, "z.a")
+
     def test_lineage_cte_union(self) -> None:
         query = """
         WITH dataset AS (