From 20739a12c39121a9e7ad3c9a2469ec5a6876199d Mon Sep 17 00:00:00 2001
From: Daniel Baumann <daniel.baumann@progress-linux.org>
Date: Sat, 3 Jun 2023 01:59:40 +0200
Subject: Merging upstream version 15.0.0.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
---
 tests/test_lineage.py | 46 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 7 deletions(-)

(limited to 'tests/test_lineage.py')

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, [])
-- 
cgit v1.2.3