1 files changed, 39 insertions, 7 deletions
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index 1d13dd3..f33a2c2 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -2,12 +2,18 @@ from __future__ import annotations
 
 import unittest
 
+import sqlglot
 from sqlglot.lineage import lineage
+from sqlglot.schema import MappingSchema
 
 
 class TestLineage(unittest.TestCase):
     maxDiff = None
 
+    @classmethod
+    def setUpClass(cls):
+        sqlglot.schema = MappingSchema()
+
     def test_lineage(self) -> None:
         node = lineage(
             "a",
@@ -140,17 +146,43 @@ class TestLineage(unittest.TestCase):
         self.assertEqual(node.alias, "")
 
         downstream = node.downstream[0]
-        self.assertEqual(
-            downstream.source.sql(),
-            "SELECT t.a AS a FROM (VALUES (1), (2)) AS t(a)",
-        )
+        self.assertEqual(downstream.source.sql(), "SELECT t.a AS a FROM (VALUES (1), (2)) AS t(a)")
         self.assertEqual(downstream.expression.sql(), "t.a AS a")
         self.assertEqual(downstream.alias, "y")
 
         downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "(VALUES (1), (2)) AS t(a)")
+        self.assertEqual(downstream.expression.sql(), "a")
+        self.assertEqual(downstream.alias, "")
+
+    def test_lineage_cte_name_appears_in_schema(self) -> None:
+        schema = {"a": {"b": {"t1": {"c1": "int"}, "t2": {"c2": "int"}}}}
+
+        node = lineage(
+            "c2",
+            "WITH t1 AS (SELECT * FROM a.b.t2), inter AS (SELECT * FROM t1) SELECT * FROM inter",
+            schema=schema,
+        )
+
         self.assertEqual(
-            downstream.source.sql(),
-            "(VALUES (1), (2)) AS t(a)",
+            node.source.sql(),
+            "WITH t1 AS (SELECT t2.c2 AS c2 FROM a.b.t2 AS t2), inter AS (SELECT t1.c2 AS c2 FROM t1) SELECT inter.c2 AS c2 FROM inter",
         )
-        self.assertEqual(downstream.expression.sql(), "a")
+        self.assertEqual(node.alias, "")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.source.sql(), "SELECT t1.c2 AS c2 FROM t1")
+        self.assertEqual(downstream.expression.sql(), "t1.c2 AS c2")
+        self.assertEqual(downstream.alias, "")
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "SELECT t2.c2 AS c2 FROM a.b.t2 AS t2")
+        self.assertEqual(downstream.expression.sql(), "t2.c2 AS c2")
+        self.assertEqual(downstream.alias, "")
+
+        downstream = downstream.downstream[0]
+        self.assertEqual(downstream.source.sql(), "a.b.t2 AS t2")
+        self.assertEqual(downstream.expression.sql(), "a.b.t2 AS t2")
         self.assertEqual(downstream.alias, "")
+
+        self.assertEqual(downstream.downstream, [])