Add unit test for converson with source code

Fix _convert_lineage_source when converting advanced
custom lineage v1. Was not accurately converting code
highlights.

tests/test_data/conversion/lineage.json

@@ -105,8 +105,27 @@
             "column": "col2"
           }
         ],
-        "mapping": "VIEW1 creation",
-        "source_code": "SELECT col1, col2 from T1;"
-      }]
+        "mapping_ref": {
+          "mapping": "VIEW1_creation",
+          "source_code": "transforms.sql",
+          "codebase_pos": [
+            {
+              "pos_start": 0,
+              "pos_len": 26
+            }
+          ]
+        }
+      }
+    ],
+    "codebase_files": {
+      "transforms.sql": {
+        "mapping_refs": {
+          "VIEW1_creation": {
+            "pos_start": 0,
+            "pos_len": 26
+          }
+        }
+      }
+    }
   }
 

tests/test_data/conversion/lineage_v3.json

@@ -47,6 +47,10 @@
         "name": "col1",
         "type": "Column"
       }
+    },
+    "source_code": {
+      "path": "source_codes/uuid.txt",
+      "transformation_display_name": "VIEW1 creation"
     }
   },
   {
@@ -97,6 +101,16 @@
         "name": "col2",
         "type": "Column"
       }
+    },
+    "source_code": {
+      "path": "source_codes/uuid.txt",
+      "highlights": [
+        {
+          "start": 0,
+          "len": 26
+        }
+      ],
+      "transformation_display_name": "VIEW1_creation"
     }
   }
 ]

tests/test_data/conversion/lineage_v3_no_source_code.json

@@ -0,0 +1,102 @@
+[
+  {
+    "src": {
+      "nodes": [
+        {
+          "name": "snowflake",
+          "type": "System"
+        },
+        {
+          "name": "DB1",
+          "type": "Database"
+        },
+        {
+          "name": "PUBLIC",
+          "type": "Schema"
+        }
+      ],
+      "parent": {
+        "name": "T1",
+        "type": "Table"
+      },
+      "leaf": {
+        "name": "col1",
+        "type": "Column"
+      }
+    },
+    "trg": {
+      "nodes": [
+        {
+          "name": "snowflake",
+          "type": "System"
+        },
+        {
+          "name": "DB1",
+          "type": "Database"
+        },
+        {
+          "name": "PUBLIC",
+          "type": "Schema"
+        }
+      ],
+      "parent": {
+        "name": "VIEW1",
+        "type": "Table"
+      },
+      "leaf": {
+        "name": "col1",
+        "type": "Column"
+      }
+    }
+  },
+  {
+    "src": {
+      "nodes": [
+        {
+          "name": "snowflake",
+          "type": "System"
+        },
+        {
+          "name": "DB1",
+          "type": "Database"
+        },
+        {
+          "name": "PUBLIC",
+          "type": "Schema"
+        }
+      ],
+      "parent": {
+        "name": "T1",
+        "type": "Table"
+      },
+      "leaf": {
+        "name": "col2",
+        "type": "Column"
+      }
+    },
+    "trg": {
+      "nodes": [
+        {
+          "name": "snowflake",
+          "type": "System"
+        },
+        {
+          "name": "DB1",
+          "type": "Database"
+        },
+        {
+          "name": "PUBLIC",
+          "type": "Schema"
+        }
+      ],
+      "parent": {
+        "name": "VIEW1",
+        "type": "Table"
+      },
+      "leaf": {
+        "name": "col2",
+        "type": "Column"
+      }
+    }
+  }
+]

tests/test_data/conversion/transforms.sql

@@ -0,0 +1 @@
+SELECT col1, col2 from T1;

tests/test_translate_to_batch_format.py

@@ -4,6 +4,37 @@
 from tools.translate_to_batch_format import convert
 
 
+def test_translate_with_simple_and_advanced_source_code() -> None:
+    convert(
+        input_directory="./test_data/conversion",
+        output_directory="./test_data/conversion/v3",
+        migrate_source_code=True,
+    )
+
+    # compare converted with expected
+    with open("./test_data/conversion/metadata.json") as input_file:
+        expected_metadata = json.load(input_file)
+
+    with open("./test_data/conversion/v3/metadata.json") as input_file:
+        generated_metadata = json.load(input_file)
+
+    with open("./test_data/conversion/lineage_v3.json") as input_file:
+        expected_lineage = json.load(input_file)
+
+    with open("./test_data/conversion/v3/lineage.json") as input_file:
+        generated_lineage = json.load(input_file)
+
+    for lineage in generated_lineage:
+        if lineage.get("source_code"):
+            lineage["source_code"]["path"] = "source_codes/uuid.txt"
+
+    assert expected_metadata == generated_metadata
+    assert expected_lineage == generated_lineage
+
+    # cleanup
+    shutil.rmtree("./test_data/conversion/v3", ignore_errors=True)
+
+
 def test_translat_without_source_code() -> None:
     # convert input
     convert(
@@ -19,7 +50,7 @@ def test_translat_without_source_code() -> None:
     with open("./test_data/conversion/v3/metadata.json") as input_file:
         generated_metadata = json.load(input_file)
 
-    with open("./test_data/conversion/lineage_v3.json") as input_file:
+    with open("./test_data/conversion/lineage_v3_no_source_code.json") as input_file:
         expected_lineage = json.load(input_file)
 
     with open("./test_data/conversion/v3/lineage.json") as input_file:

tools/translate_to_batch_format.py

@@ -67,7 +67,7 @@ def _convert_lineage_source(
     if "mapping_ref" in lineage_relationship_v1:
         source_code_file_v1 = lineage_relationship_v1.get("mapping_ref", {}).get("source_code", "")
         mapping_v1 = lineage_relationship_v1.get("mapping_ref", {}).get("mapping", "")
-        codebase_pos_v1 = lineage_relationship_v1.get("codebase_pos", [])
+        codebase_pos_v1 = lineage_relationship_v1.get("mapping_ref", {}).get("codebase_pos", [])
         if not source_code_file_v1:
             return None
 
@@ -88,7 +88,7 @@ def _convert_lineage_source(
                     custom_lineage_config=custom_lineage_config,
                     transformation_display_name=mapping_v1,
                     highlights=[
-                        SourceCodeHighLight(start=highlight_v1["post_start"], len=highlight_v1["pos_len"])
+                        SourceCodeHighLight(start=highlight_v1["pos_start"], len=highlight_v1["pos_len"])
                         for highlight_v1 in codebase_pos_v1
                     ],
                 )