wehub-resource-sync
bf2343b7e4
Integration Tests - MySQL + Elasticsearch / Detect Changes (push) Has been cancelled
Integration Tests - MySQL + Elasticsearch / integration-tests-mysql-elasticsearch (push) Has been cancelled
Integration Tests - PostgreSQL + Elasticsearch + Redis / Detect Changes (push) Has been cancelled
Integration Tests - PostgreSQL + Elasticsearch + Redis / integration-tests-postgres-elasticsearch-redis (push) Has been cancelled
Integration Tests - PostgreSQL + OpenSearch / Detect Changes (push) Has been cancelled
Integration Tests - PostgreSQL + OpenSearch / integration-tests-postgres-opensearch (push) Has been cancelled
Java Checkstyle / java-checkstyle (push) Has been cancelled
Maven Collate Tests / maven-collate-ci (push) Has been cancelled
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests-status (push) Has been cancelled
Publish Package to Maven Central Repository / publish-maven-packages (push) Has been cancelled
OpenMetadata Service Unit Tests / Detect Changes (push) Has been cancelled
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests (push) Has been cancelled
OpenMetadata Service Unit Tests / k8s_operator-unit-tests (push) Has been cancelled
1226 lines
50 KiB
Python
1226 lines
50 KiB
Python
# Copyright 2025 Collate
|
|
# Licensed under the Collate Community License, Version 1.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Test datalake utils
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from unittest import TestCase
|
|
|
|
import pandas as pd
|
|
|
|
from metadata.generated.schema.entity.data.table import Column, DataType
|
|
from metadata.readers.dataframe.dsv import DSVDataFrameReader
|
|
from metadata.readers.dataframe.reader_factory import SupportedTypes
|
|
from metadata.utils.datalake.datalake_utils import (
|
|
DataFrameColumnParser,
|
|
GenericDataFrameColumnParser,
|
|
JsonDataFrameColumnParser,
|
|
ParquetDataFrameColumnParser,
|
|
get_file_format_type,
|
|
)
|
|
|
|
STRUCTURE = {
|
|
"a": "w",
|
|
"b": 4,
|
|
"c": {
|
|
"d": 2,
|
|
"e": 4,
|
|
"f": {
|
|
"g": 9,
|
|
"h": {"i": 6},
|
|
"n": {
|
|
"o": 10,
|
|
"p": 11,
|
|
},
|
|
},
|
|
"j": 7,
|
|
"k": 8,
|
|
},
|
|
}
|
|
|
|
|
|
class TestDatalakeUtils(TestCase):
|
|
"""class for datalake utils test"""
|
|
|
|
def test_unique_json_structure(self):
|
|
"""test unique json structure fn"""
|
|
sample_data = [
|
|
{"a": "x", "b": 1, "c": {"d": 2}},
|
|
{"a": "y", "b": 2, "c": {"e": 4, "f": {"g": 5, "h": {"i": 6}, "n": 5}}},
|
|
{"a": "z", "b": 3, "c": {"j": 7}},
|
|
{"a": "w", "b": 4, "c": {"k": 8, "f": {"g": 9, "n": {"o": 10, "p": 11}}}},
|
|
]
|
|
expected = STRUCTURE
|
|
|
|
actual = GenericDataFrameColumnParser.unique_json_structure(sample_data)
|
|
|
|
self.assertDictEqual(expected, actual)
|
|
|
|
def test_construct_column(self):
|
|
"""test construct column fn"""
|
|
expected = [
|
|
{
|
|
"dataTypeDisplay": "STRING",
|
|
"dataType": "STRING",
|
|
"name": "a",
|
|
"displayName": "a",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "b",
|
|
"displayName": "b",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "JSON",
|
|
"dataType": "JSON",
|
|
"name": "c",
|
|
"displayName": "c",
|
|
"children": [
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "d",
|
|
"displayName": "d",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "e",
|
|
"displayName": "e",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "JSON",
|
|
"dataType": "JSON",
|
|
"name": "f",
|
|
"displayName": "f",
|
|
"children": [
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "g",
|
|
"displayName": "g",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "JSON",
|
|
"dataType": "JSON",
|
|
"name": "h",
|
|
"displayName": "h",
|
|
"children": [
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "i",
|
|
"displayName": "i",
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"dataTypeDisplay": "JSON",
|
|
"dataType": "JSON",
|
|
"name": "n",
|
|
"displayName": "n",
|
|
"children": [
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "o",
|
|
"displayName": "o",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "p",
|
|
"displayName": "p",
|
|
},
|
|
],
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "j",
|
|
"displayName": "j",
|
|
},
|
|
{
|
|
"dataTypeDisplay": "INT",
|
|
"dataType": "INT",
|
|
"name": "k",
|
|
"displayName": "k",
|
|
},
|
|
],
|
|
},
|
|
]
|
|
actual = GenericDataFrameColumnParser.construct_json_column_children(STRUCTURE)
|
|
|
|
for el in zip(expected, actual): # noqa: B905
|
|
self.assertDictEqual(el[0], el[1])
|
|
|
|
def test_unique_json_structure_with_list_of_dicts(self):
|
|
"""list-of-dicts values are merged into a struct shape (e.g. Iceberg `schema.fields`)."""
|
|
sample_data = [
|
|
{
|
|
"schema": {
|
|
"fields": [
|
|
{"id": 1, "name": "customer_id", "type": "string"},
|
|
{"id": 2, "name": "customer_type_cd", "type": "string"},
|
|
]
|
|
}
|
|
}
|
|
]
|
|
|
|
actual = GenericDataFrameColumnParser.unique_json_structure(sample_data)
|
|
fields_value = actual["schema"]["fields"]
|
|
|
|
from metadata.utils.datalake.datalake_utils import _ArrayOfStruct
|
|
|
|
assert isinstance(fields_value, _ArrayOfStruct)
|
|
assert set(fields_value.struct.keys()) == {"id", "name", "type"}
|
|
|
|
def test_unique_json_structure_merges_list_of_dicts_across_samples(self):
|
|
"""list-of-dicts values across multiple samples are unioned, not overwritten."""
|
|
from metadata.utils.datalake.datalake_utils import _ArrayOfStruct
|
|
|
|
sample_data = [
|
|
{"schema": {"fields": [{"id": 1, "name": "customer_id", "type": "string"}]}},
|
|
{"schema": {"fields": [{"id": 2, "required": False, "type": "string"}]}},
|
|
{"schema": {"fields": [{"description": "ciam id"}]}},
|
|
]
|
|
|
|
actual = GenericDataFrameColumnParser.unique_json_structure(sample_data)
|
|
fields_value = actual["schema"]["fields"]
|
|
|
|
assert isinstance(fields_value, _ArrayOfStruct)
|
|
assert set(fields_value.struct.keys()) == {"id", "name", "type", "required", "description"}
|
|
|
|
def test_construct_column_with_array_of_struct(self):
|
|
"""list-of-dicts values render as ARRAY<STRUCT<...>> with children for the struct fields."""
|
|
structure = {
|
|
"schema": {
|
|
"fields": [
|
|
{"id": 1, "name": "customer_id", "type": "string"},
|
|
{"id": 2, "name": "ciam_id", "type": "string"},
|
|
]
|
|
}
|
|
}
|
|
merged = GenericDataFrameColumnParser.unique_json_structure([structure])
|
|
children = GenericDataFrameColumnParser.construct_json_column_children(merged)
|
|
|
|
schema_col = children[0]
|
|
fields_col = next(c for c in schema_col["children"] if c["name"] == "fields")
|
|
|
|
assert fields_col["dataType"] == DataType.ARRAY.value
|
|
assert fields_col["arrayDataType"] == DataType.STRUCT
|
|
assert {child["name"] for child in fields_col["children"]} == {"id", "name", "type"}
|
|
|
|
def test_create_column_object(self):
|
|
"""test create column object fn"""
|
|
formatted_column = GenericDataFrameColumnParser.construct_json_column_children(STRUCTURE)
|
|
column = {
|
|
"dataTypeDisplay": "STRING",
|
|
"dataType": "STRING",
|
|
"name": "a",
|
|
"displayName": "a",
|
|
"children": formatted_column,
|
|
}
|
|
column_obj = Column(**column)
|
|
assert column_obj.children is not None and len(column_obj.children) == 3
|
|
|
|
def test_fetch_col_types_majority_wins(self):
|
|
"""Majority type wins; a handful of date-parseable tokens must not flip a string column."""
|
|
cases = [
|
|
# Overwhelmingly strings with a few month-name values — must stay STRING.
|
|
# This is the dvdrental last_name bug: "May" parses as a date via dateutil
|
|
# but the column is a string column.
|
|
(
|
|
"last_name_with_month_surnames",
|
|
["Smith", "Gonzalez", "Brown", "May", "Jones", "Williams", "Davis"],
|
|
DataType.STRING,
|
|
),
|
|
# Minority of ambiguous month tokens mixed in a long list of plain strings.
|
|
("mostly_strings_few_month_tokens", ["foo", "bar", "baz", "May", "qux", "quux", "March"], DataType.STRING),
|
|
# All values are unambiguous ISO dates — must be DATETIME.
|
|
("pure_iso_dates", ["2024-01-01", "2024-06-15", "2025-03-20"], DataType.DATETIME),
|
|
# Natural-language date phrases — all parse as dates — must be DATETIME.
|
|
("natural_language_dates", ["May 2025", "June 2026", "March 2024", "January 2023"], DataType.DATETIME),
|
|
# Pure strings, no date-parseable values at all.
|
|
("pure_strings", ["hello", "world", "foo", "bar"], DataType.STRING),
|
|
# All plain integers stored as strings — must be INT.
|
|
("integer_strings", ["1", "2", "3", "42"], DataType.INT),
|
|
]
|
|
for name, values, expected in cases:
|
|
with self.subTest(name):
|
|
df = pd.DataFrame({"col": values})
|
|
self.assertEqual(GenericDataFrameColumnParser.fetch_col_types(df, "col"), expected)
|
|
|
|
|
|
class TestParquetDataFrameColumnParser(TestCase):
|
|
"""Test parquet dataframe column parser"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
resources_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resources") # noqa: PTH118, PTH120
|
|
cls.parquet_path = os.path.join(resources_path, "datalake", "example.parquet") # noqa: PTH118
|
|
|
|
cls.df = pd.read_parquet(cls.parquet_path)
|
|
|
|
cls.parquet_parser = ParquetDataFrameColumnParser(cls.df)
|
|
|
|
def test_parser_instantiation(self):
|
|
"""Test the right parser is instantiated from the creator method"""
|
|
parquet_parser = DataFrameColumnParser.create(self.df, SupportedTypes.PARQUET)
|
|
self.assertIsInstance(parquet_parser.parser, ParquetDataFrameColumnParser)
|
|
|
|
parquet_types = [
|
|
SupportedTypes.PARQUET,
|
|
SupportedTypes.PARQUET_PQ,
|
|
SupportedTypes.PARQUET_PQT,
|
|
SupportedTypes.PARQUET_PARQ,
|
|
SupportedTypes.PARQUET_SNAPPY,
|
|
]
|
|
other_types = [typ for typ in SupportedTypes if typ not in parquet_types]
|
|
for other_type in other_types:
|
|
with self.subTest(other_type=other_type):
|
|
generic_parser = DataFrameColumnParser.create(self.df, other_type)
|
|
self.assertIsInstance(generic_parser.parser, GenericDataFrameColumnParser)
|
|
|
|
def test_shuffle_and_sample_from_parser(self):
|
|
"""test the shuffle and sampling logic from the parser creator method"""
|
|
parquet_parser = DataFrameColumnParser.create(self.df, SupportedTypes.PARQUET)
|
|
self.assertEqual(parquet_parser.parser.data_frame.shape, self.df.shape)
|
|
|
|
parquet_parser = DataFrameColumnParser.create([self.df, self.df], SupportedTypes.PARQUET)
|
|
self.assertEqual(parquet_parser.parser.data_frame.shape, self.df.shape)
|
|
|
|
parquet_parser = DataFrameColumnParser.create([self.df, self.df], SupportedTypes.PARQUET, sample=False)
|
|
self.assertEqual(parquet_parser.parser.data_frame.shape, pd.concat([self.df, self.df]).shape)
|
|
|
|
def test_get_columns(self):
|
|
"""test `get_columns` method of the parquet column parser"""
|
|
expected = [
|
|
Column(
|
|
dataTypeDisplay="bool",
|
|
dataType=DataType.BOOLEAN,
|
|
name="a",
|
|
displayName="a",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int8",
|
|
dataType=DataType.INT,
|
|
name="b",
|
|
displayName="b",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int16",
|
|
dataType=DataType.INT,
|
|
name="c",
|
|
displayName="c",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int32",
|
|
dataType=DataType.INT,
|
|
name="d",
|
|
displayName="d",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="e",
|
|
displayName="e",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="uint8",
|
|
dataType=DataType.UINT,
|
|
name="f",
|
|
displayName="f",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="uint16",
|
|
dataType=DataType.UINT,
|
|
name="g",
|
|
displayName="g",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="uint32",
|
|
dataType=DataType.UINT,
|
|
name="h",
|
|
displayName="h",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="uint64",
|
|
dataType=DataType.UINT,
|
|
name="i",
|
|
displayName="i",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="float",
|
|
dataType=DataType.FLOAT,
|
|
name="k",
|
|
displayName="k",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="double",
|
|
dataType=DataType.FLOAT,
|
|
name="l",
|
|
displayName="l",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="time64[us]",
|
|
dataType=DataType.DATETIME,
|
|
name="n",
|
|
displayName="n",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="timestamp[ns]",
|
|
dataType=DataType.DATETIME,
|
|
name="o",
|
|
displayName="o",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="date32[day]",
|
|
dataType=DataType.DATE,
|
|
name="p",
|
|
displayName="p",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="date32[day]",
|
|
dataType=DataType.DATE,
|
|
name="q",
|
|
displayName="q",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="duration[ns]",
|
|
dataType=DataType.INT,
|
|
name="r",
|
|
displayName="r",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="binary",
|
|
dataType=DataType.BINARY,
|
|
name="t",
|
|
displayName="t",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="string",
|
|
dataType=DataType.STRING,
|
|
name="u",
|
|
displayName="u",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="string",
|
|
dataType=DataType.STRING,
|
|
name="v",
|
|
displayName="v",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="binary",
|
|
dataType=DataType.BINARY,
|
|
name="w",
|
|
displayName="w",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="string",
|
|
dataType=DataType.STRING,
|
|
name="x",
|
|
displayName="x",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="string",
|
|
dataType=DataType.STRING,
|
|
name="y",
|
|
displayName="y",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="list<item: int64>",
|
|
dataType=DataType.ARRAY,
|
|
name="aa",
|
|
displayName="aa",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="list<item: int64>",
|
|
dataType=DataType.ARRAY,
|
|
name="bb",
|
|
displayName="bb",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="struct<ee: int64, ff: int64, gg: struct<hh: struct<ii: int64, jj: int64, kk: int64>>>",
|
|
dataType=DataType.STRUCT,
|
|
name="dd",
|
|
displayName="dd",
|
|
children=[
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="ee",
|
|
displayName="ee",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="ff",
|
|
displayName="ff",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="struct<hh: struct<ii: int64, jj: int64, kk: int64>>",
|
|
dataType=DataType.STRUCT,
|
|
name="gg",
|
|
displayName="gg",
|
|
children=[
|
|
Column(
|
|
dataTypeDisplay="struct<ii: int64, jj: int64, kk: int64>",
|
|
dataType=DataType.STRUCT,
|
|
name="hh",
|
|
displayName="hh",
|
|
children=[
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="ii",
|
|
displayName="ii",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="jj",
|
|
displayName="jj",
|
|
), # type: ignore
|
|
Column(
|
|
dataTypeDisplay="int64",
|
|
dataType=DataType.INT,
|
|
name="kk",
|
|
displayName="kk",
|
|
), # type: ignore
|
|
],
|
|
),
|
|
],
|
|
),
|
|
],
|
|
), # type: ignore
|
|
]
|
|
actual = self.parquet_parser.get_columns()
|
|
for validation in zip(expected, actual): # noqa: B905
|
|
with self.subTest(validation=validation):
|
|
expected_col, actual_col = validation
|
|
self.assertEqual(expected_col.name, actual_col.name)
|
|
self.assertEqual(expected_col.displayName, actual_col.displayName)
|
|
self.assertEqual(expected_col.dataType, actual_col.dataType)
|
|
|
|
def _validate_parsed_column(self, expected, actual):
|
|
"""validate parsed column"""
|
|
self.assertEqual(expected.name, actual.name)
|
|
self.assertEqual(expected.dataType, actual.dataType)
|
|
self.assertEqual(expected.displayName, actual.displayName)
|
|
if expected.children:
|
|
self.assertEqual(len(expected.children), len(actual.children))
|
|
for validation in zip(expected.children, actual.children): # noqa: B905
|
|
with self.subTest(validation=validation):
|
|
expected_col, actual_col = validation
|
|
self._validate_parsed_column(expected_col, actual_col)
|
|
|
|
def test_get_file_format_type_csv_gz(self):
|
|
"""test get_file_format_type function for csv.gz files"""
|
|
# Test csv.gz file detection
|
|
result = get_file_format_type("data.csv.gz")
|
|
self.assertEqual(result, SupportedTypes.CSVGZ)
|
|
|
|
# Test regular csv file detection (should still work)
|
|
result = get_file_format_type("data.csv")
|
|
self.assertEqual(result, SupportedTypes.CSV)
|
|
|
|
# Test other gzipped files
|
|
result = get_file_format_type("data.json.gz")
|
|
self.assertEqual(result, SupportedTypes.JSONGZ)
|
|
|
|
# Test unsupported gzipped format
|
|
result = get_file_format_type("data.txt.gz")
|
|
self.assertEqual(result, False)
|
|
|
|
def test_csv_gz_file_format_detection_edge_cases(self):
|
|
"""test edge cases for csv.gz file format detection"""
|
|
# Test with nested paths
|
|
result = get_file_format_type("folder/subfolder/data.csv.gz")
|
|
self.assertEqual(result, SupportedTypes.CSVGZ)
|
|
|
|
# Test with multiple dots
|
|
result = get_file_format_type("data.backup.csv.gz")
|
|
self.assertEqual(result, SupportedTypes.CSVGZ)
|
|
|
|
# Test with no extension
|
|
result = get_file_format_type("data")
|
|
self.assertEqual(result, False)
|
|
|
|
# Test with just .gz
|
|
result = get_file_format_type("data.gz")
|
|
self.assertEqual(result, False)
|
|
|
|
def test_csv_gz_compression_detection(self):
|
|
"""test compression detection for various file types"""
|
|
# Test csv.gz compression detection
|
|
test_cases = [
|
|
("data.csv.gz", SupportedTypes.CSVGZ),
|
|
("data.csv", SupportedTypes.CSV),
|
|
("data.json.gz", SupportedTypes.JSONGZ),
|
|
("data.json", SupportedTypes.JSON),
|
|
("data.jsonl.gz", SupportedTypes.JSONLGZ),
|
|
("data.jsonl", SupportedTypes.JSONL),
|
|
("data.parquet", SupportedTypes.PARQUET),
|
|
("data.txt.gz", False), # Unsupported
|
|
("data.unknown.gz", False), # Unsupported
|
|
]
|
|
|
|
for filename, expected in test_cases:
|
|
with self.subTest(filename=filename):
|
|
result = get_file_format_type(filename)
|
|
self.assertEqual(result, expected, f"Failed for {filename}")
|
|
|
|
def test_csv_gz_reader_factory_integration(self):
|
|
"""test that csv.gz is properly integrated with reader factory"""
|
|
from metadata.readers.dataframe.reader_factory import SupportedTypes
|
|
|
|
# Test that CSVGZ is properly handled
|
|
try:
|
|
# Test that the enum value exists
|
|
self.assertEqual(SupportedTypes.CSVGZ.value, "csv.gz")
|
|
|
|
# Test that it's different from regular CSV
|
|
self.assertNotEqual(SupportedTypes.CSVGZ, SupportedTypes.CSV)
|
|
self.assertNotEqual(SupportedTypes.CSVGZ.value, SupportedTypes.CSV.value)
|
|
|
|
except Exception as e:
|
|
self.fail(f"CSVGZ enum test failed: {e}")
|
|
|
|
def test_csv_gz_supported_types_enum(self):
|
|
"""test that CSVGZ is properly defined in SupportedTypes enum"""
|
|
# Test that CSVGZ exists in the enum
|
|
self.assertIn(SupportedTypes.CSVGZ, SupportedTypes)
|
|
self.assertEqual(SupportedTypes.CSVGZ.value, "csv.gz")
|
|
|
|
# Test that it's different from regular CSV
|
|
self.assertNotEqual(SupportedTypes.CSVGZ, SupportedTypes.CSV)
|
|
self.assertNotEqual(SupportedTypes.CSVGZ.value, SupportedTypes.CSV.value)
|
|
|
|
def test_csv_gz_dsv_reader_compression_detection(self):
|
|
"""test that DSV reader properly detects compression for csv.gz files"""
|
|
from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
|
|
LocalConfig,
|
|
)
|
|
from metadata.readers.dataframe.dsv import DSVDataFrameReader
|
|
|
|
# Create a mock config
|
|
local_config = LocalConfig()
|
|
|
|
# Create DSV reader
|
|
reader = DSVDataFrameReader(config_source=local_config, client=None) # noqa: F841
|
|
|
|
# Test compression detection logic (this is the same logic used in the dispatch methods)
|
|
test_cases = [
|
|
("data.csv.gz", "gzip"),
|
|
("data.csv", None),
|
|
("data.json.gz", "gzip"),
|
|
("data.txt.gz", "gzip"),
|
|
("data.unknown.gz", "gzip"),
|
|
]
|
|
|
|
for filename, expected_compression in test_cases:
|
|
with self.subTest(filename=filename):
|
|
# Simulate the compression detection logic from the dispatch methods
|
|
compression = None
|
|
if filename.endswith(".gz"):
|
|
compression = "gzip"
|
|
|
|
self.assertEqual(
|
|
compression,
|
|
expected_compression,
|
|
f"Compression detection failed for {filename}",
|
|
)
|
|
|
|
def test_csv_gz_integration_completeness(self):
|
|
"""test that csv.gz support is complete across all components"""
|
|
# Test that CSVGZ is in the reader factory mapping
|
|
from metadata.readers.dataframe.reader_factory import (
|
|
DF_READER_MAP,
|
|
SupportedTypes,
|
|
)
|
|
|
|
# Check that CSVGZ is mapped to CSVDataFrameReader
|
|
self.assertIn(SupportedTypes.CSVGZ.value, DF_READER_MAP)
|
|
|
|
# Test that the get_df_reader function includes CSVGZ in DSV handling
|
|
|
|
# This should not raise an exception for CSVGZ
|
|
try:
|
|
# Test that CSVGZ is included in the DSV types
|
|
dsv_types = {SupportedTypes.CSV, SupportedTypes.CSVGZ, SupportedTypes.TSV}
|
|
self.assertIn(SupportedTypes.CSVGZ, dsv_types)
|
|
except Exception as e:
|
|
self.fail(f"CSVGZ integration test failed: {e}")
|
|
|
|
|
|
class TestIcebergDeltaLakeMetadataParsing(TestCase):
|
|
"""Test Iceberg/Delta Lake metadata JSON parsing"""
|
|
|
|
def test_iceberg_metadata_parsing(self):
|
|
"""Test parsing of Iceberg/Delta Lake metadata files with nested schema.fields structure"""
|
|
|
|
# Sample Iceberg/Delta Lake metadata structure
|
|
iceberg_metadata = {
|
|
"format-version": 1,
|
|
"table-uuid": "e9182d72-131b-48fe-b530-79edc044fb01",
|
|
"location": "s3://bucket/path/table",
|
|
"schema": {
|
|
"type": "struct",
|
|
"schema-id": 0,
|
|
"fields": [
|
|
{
|
|
"id": 1,
|
|
"name": "customer_id",
|
|
"required": False,
|
|
"type": "string",
|
|
},
|
|
{
|
|
"id": 2,
|
|
"name": "customer_type_cd",
|
|
"required": False,
|
|
"type": "string",
|
|
},
|
|
{"id": 3, "name": "amount", "required": True, "type": "double"},
|
|
{
|
|
"id": 4,
|
|
"name": "is_active",
|
|
"required": False,
|
|
"type": "boolean",
|
|
},
|
|
{"id": 5, "name": "order_count", "required": False, "type": "int"},
|
|
{
|
|
"id": 6,
|
|
"name": "created_date",
|
|
"required": False,
|
|
"type": "date",
|
|
},
|
|
{
|
|
"id": 7,
|
|
"name": "updated_timestamp",
|
|
"required": False,
|
|
"type": "timestamp",
|
|
},
|
|
{
|
|
"id": 8,
|
|
"name": "metadata",
|
|
"required": False,
|
|
"type": {
|
|
"type": "struct",
|
|
"fields": [
|
|
{
|
|
"id": 9,
|
|
"name": "source_system",
|
|
"required": False,
|
|
"type": "string",
|
|
},
|
|
{
|
|
"id": 10,
|
|
"name": "last_sync_time",
|
|
"required": False,
|
|
"type": "timestamp",
|
|
},
|
|
],
|
|
},
|
|
},
|
|
],
|
|
},
|
|
}
|
|
|
|
# Convert to JSON string as would be received from file
|
|
raw_data = json.dumps(iceberg_metadata)
|
|
|
|
# Create a dummy DataFrame (required by parser but not used for Iceberg metadata)
|
|
df = pd.DataFrame()
|
|
|
|
# Create parser and parse columns
|
|
parser = JsonDataFrameColumnParser(df, raw_data=raw_data)
|
|
columns = parser.get_columns()
|
|
|
|
# Verify the correct number of columns were parsed
|
|
self.assertEqual(len(columns), 8)
|
|
|
|
# Verify field names were correctly parsed
|
|
expected_names = [
|
|
"customer_id",
|
|
"customer_type_cd",
|
|
"amount",
|
|
"is_active",
|
|
"order_count",
|
|
"created_date",
|
|
"updated_timestamp",
|
|
"metadata",
|
|
]
|
|
actual_names = [col.displayName for col in columns]
|
|
self.assertEqual(expected_names, actual_names)
|
|
|
|
# Verify data types were correctly mapped
|
|
expected_types = [
|
|
DataType.STRING, # customer_id
|
|
DataType.STRING, # customer_type_cd
|
|
DataType.DOUBLE, # amount
|
|
DataType.BOOLEAN, # is_active
|
|
DataType.INT, # order_count
|
|
DataType.DATE, # created_date
|
|
DataType.TIMESTAMP, # updated_timestamp
|
|
DataType.STRUCT, # metadata
|
|
]
|
|
actual_types = [col.dataType for col in columns]
|
|
self.assertEqual(expected_types, actual_types)
|
|
|
|
# Verify nested struct field (metadata)
|
|
metadata_column = columns[7]
|
|
self.assertEqual(metadata_column.displayName, "metadata")
|
|
self.assertEqual(metadata_column.dataType, DataType.STRUCT)
|
|
self.assertIsNotNone(metadata_column.children)
|
|
self.assertEqual(len(metadata_column.children), 2)
|
|
|
|
# Verify nested field details
|
|
nested_fields = metadata_column.children
|
|
self.assertEqual(nested_fields[0]["displayName"], "source_system")
|
|
self.assertEqual(nested_fields[0]["dataType"], DataType.STRING.value)
|
|
self.assertEqual(nested_fields[1]["displayName"], "last_sync_time")
|
|
self.assertEqual(nested_fields[1]["dataType"], DataType.TIMESTAMP.value)
|
|
|
|
def test_is_iceberg_delta_metadata_detection(self):
|
|
"""Test detection of Iceberg/Delta Lake metadata format"""
|
|
df = pd.DataFrame()
|
|
parser = JsonDataFrameColumnParser(df, raw_data=None)
|
|
|
|
# Test valid Iceberg/Delta Lake metadata
|
|
valid_metadata = {"schema": {"fields": [{"name": "field1", "type": "string"}]}}
|
|
self.assertTrue(parser._is_iceberg_delta_metadata(valid_metadata))
|
|
|
|
# Test invalid formats
|
|
invalid_cases = [
|
|
{}, # Empty dict
|
|
{"schema": "not_a_dict"}, # Schema not a dict
|
|
{"schema": {}}, # No fields
|
|
{"schema": {"fields": "not_a_list"}}, # Fields not a list
|
|
{"properties": {}}, # JSON Schema format (not Iceberg)
|
|
]
|
|
|
|
for invalid_case in invalid_cases:
|
|
with self.subTest(invalid_case=invalid_case):
|
|
self.assertFalse(parser._is_iceberg_delta_metadata(invalid_case))
|
|
|
|
def test_fallback_to_json_schema_parser(self):
|
|
"""Test that non-Iceberg JSON files fall back to standard JSON Schema parser"""
|
|
# Standard JSON Schema format
|
|
json_schema = {
|
|
"$schema": "http://json-schema.org/draft-07/schema#",
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
|
|
}
|
|
|
|
raw_data = json.dumps(json_schema)
|
|
df = pd.DataFrame()
|
|
|
|
# This should use the standard JSON Schema parser, not Iceberg parser
|
|
parser = JsonDataFrameColumnParser(df, raw_data=raw_data)
|
|
columns = parser.get_columns()
|
|
|
|
# The standard parser behavior would be different
|
|
# This test ensures we don't break existing JSON Schema parsing
|
|
self.assertIsNotNone(columns)
|
|
|
|
def test_read_json_object_propagates_raw_data_for_iceberg(self):
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
content = json.dumps(
|
|
{
|
|
"format-version": 1,
|
|
"schema": {"fields": [{"id": 1, "name": "customer_id", "required": False, "type": "string"}]},
|
|
}
|
|
)
|
|
_gen, raw = JSONDataFrameReader._read_json_object(content.encode("utf-8"))
|
|
self.assertEqual(raw, content)
|
|
|
|
def test_read_json_object_returns_none_for_plain_object(self):
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
content = json.dumps({"a": 1, "b": 2})
|
|
_gen, raw = JSONDataFrameReader._read_json_object(content.encode("utf-8"))
|
|
self.assertIsNone(raw)
|
|
|
|
def test_read_json_object_propagates_raw_data_for_json_schema(self):
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
content = json.dumps({"$schema": "http://json-schema.org/draft-07/schema#", "type": "object"})
|
|
_gen, raw = JSONDataFrameReader._read_json_object(content.encode("utf-8"))
|
|
self.assertEqual(raw, content)
|
|
|
|
def test_is_json_lines_returns_false_for_minified_iceberg(self):
|
|
"""Single-line (minified) Iceberg metadata.json must NOT be treated as JSON Lines,
|
|
otherwise it bypasses _read_json_object and raw_data is never set."""
|
|
import io
|
|
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
minified = json.dumps(
|
|
{
|
|
"format-version": 1,
|
|
"schema": {"fields": [{"id": 1, "name": "customer_id", "required": False, "type": "string"}]},
|
|
}
|
|
)
|
|
self.assertFalse(JSONDataFrameReader._is_json_lines(io.BytesIO(minified.encode("utf-8"))))
|
|
|
|
def test_is_json_lines_returns_false_for_minified_json_schema(self):
|
|
import io
|
|
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
minified = json.dumps({"$schema": "http://json-schema.org/draft-07/schema#", "type": "object"})
|
|
self.assertFalse(JSONDataFrameReader._is_json_lines(io.BytesIO(minified.encode("utf-8"))))
|
|
|
|
def test_is_json_lines_returns_true_for_real_jsonl(self):
|
|
import io
|
|
|
|
from metadata.readers.dataframe.json import JSONDataFrameReader
|
|
|
|
jsonl = b'{"a": 1, "b": 2}\n{"a": 3, "b": 4}\n'
|
|
self.assertTrue(JSONDataFrameReader._is_json_lines(io.BytesIO(jsonl)))
|
|
|
|
|
|
class TestFetchColTypesWithParsedObjects:
|
|
"""fetch_col_types must correctly type object-dtype columns whose values are already
|
|
parsed Python dicts or lists, including falsy containers ({}, [])."""
|
|
|
|
def test_empty_dict_typed_as_json(self):
|
|
df = pd.DataFrame({"col": [{}]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_empty_list_typed_as_array(self):
|
|
df = pd.DataFrame({"col": [[]]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.ARRAY
|
|
|
|
def test_multiple_empty_dicts_typed_as_json(self):
|
|
df = pd.DataFrame({"col": [{}, {}, {}]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_dict_with_data_typed_as_json(self):
|
|
df = pd.DataFrame({"col": [{"k": "v"}]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_list_with_data_typed_as_array(self):
|
|
df = pd.DataFrame({"col": [[1, 2, 3]]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.ARRAY
|
|
|
|
def test_large_already_parsed_dict_typed_as_json(self):
|
|
large = {str(i): i for i in range(500)}
|
|
df = pd.DataFrame({"col": [large]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_null_column_typed_as_string(self):
|
|
df = pd.DataFrame({"col": [None]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.STRING
|
|
|
|
def test_string_column_typed_as_string(self):
|
|
df = pd.DataFrame({"col": ["hello"]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.STRING
|
|
|
|
def test_int_column_typed_as_int(self):
|
|
df = pd.DataFrame({"col": [42]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.INT
|
|
|
|
|
|
class TestFetchColTypesMixedTypes:
|
|
"""fetch_col_types must resolve the dominant type via explicit precedence, not
|
|
lexicographic max(). The old max() would return 'str' whenever a string value appeared
|
|
in the column because 'str' > 'dict' and 'str' > 'list' lexicographically."""
|
|
|
|
def test_dict_and_string_mix_typed_as_json(self):
|
|
# Previously: max(["dict", "str"]) == "str" → STRING (wrong)
|
|
# Now: precedence picks "dict" → JSON (correct)
|
|
df = pd.DataFrame({"col": [{"a": 1}, "fallback_string"]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_list_and_string_mix_typed_as_array(self):
|
|
# Previously: max(["list", "str"]) == "str" → STRING (wrong)
|
|
# Now: precedence picks "list" → ARRAY (correct)
|
|
df = pd.DataFrame({"col": [[1, 2], "fallback_string"]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.ARRAY
|
|
|
|
def test_int_and_float_mix_typed_as_float(self):
|
|
# float64 beats int64 in precedence — a column with mixed numeric types resolves to FLOAT
|
|
df = pd.DataFrame({"col": ["42", "3.14"]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.FLOAT
|
|
|
|
def test_pure_string_column_typed_as_string(self):
|
|
# Control: no structured types present → still STRING
|
|
df = pd.DataFrame({"col": ["hello", "world"]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.STRING
|
|
|
|
def test_pure_dict_column_typed_as_json(self):
|
|
# Control: all dicts → JSON with no ambiguity
|
|
df = pd.DataFrame({"col": [{"a": 1}, {"b": 2}]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
def test_dict_beats_list_in_mixed_column(self):
|
|
# dict > list in precedence
|
|
df = pd.DataFrame({"col": [{"a": 1}, [1, 2]]})
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "col") == DataType.JSON
|
|
|
|
|
|
class TestGetChildrenWithParsedDicts:
|
|
"""get_children must correctly extract children regardless of whether the Series
|
|
values are already-parsed Python dicts, JSON strings, or a mix of both."""
|
|
|
|
def test_already_parsed_dict_returns_children(self):
|
|
col = pd.Series([{"name": "Alice", "age": 30}])
|
|
children = GenericDataFrameColumnParser.get_children(col)
|
|
assert {c["name"] for c in children} == {"name", "age"}
|
|
|
|
def test_empty_dict_returns_no_children(self):
|
|
col = pd.Series([{}])
|
|
assert GenericDataFrameColumnParser.get_children(col) == []
|
|
|
|
def test_all_null_returns_no_children(self):
|
|
col = pd.Series([None, None])
|
|
assert GenericDataFrameColumnParser.get_children(col) == []
|
|
|
|
def test_string_json_returns_children(self):
|
|
col = pd.Series(['{"name": "Bob", "score": 99}'])
|
|
children = GenericDataFrameColumnParser.get_children(col)
|
|
assert {c["name"] for c in children} == {"name", "score"}
|
|
|
|
def test_mixed_string_and_dict_values_returns_union_of_children(self):
|
|
col = pd.Series(['{"a": 1, "b": 2}', {"b": 2, "c": 3}])
|
|
children = GenericDataFrameColumnParser.get_children(col)
|
|
assert {c["name"] for c in children} == {"a", "b", "c"}
|
|
|
|
def test_malformed_string_values_are_skipped(self):
|
|
col = pd.Series(["not-json", {"key": "val"}])
|
|
children = GenericDataFrameColumnParser.get_children(col)
|
|
assert {c["name"] for c in children} == {"key"}
|
|
|
|
def test_nested_dict_structure_returns_children(self):
|
|
nodes = {"model.Project.my_model": {"name": "my_model", "unique_id": "x", "description": "test"}}
|
|
col = pd.Series([nodes])
|
|
children = GenericDataFrameColumnParser.get_children(col)
|
|
assert len(children) == 1
|
|
assert children[0]["name"] == "model.Project.my_model"
|
|
|
|
|
|
class TestSingleObjectJsonFileIngestion:
|
|
"""End-to-end column parsing for single-object JSON files.
|
|
|
|
Reads fixture files with json.loads → DataFrame.from_records → _get_columns → Column objects.
|
|
A single top-level JSON object is wrapped into a 1-row DataFrame. Every top-level key
|
|
becomes a column whose value is the Python object returned by json.loads — typically a
|
|
dict, list, or None. All columns must be typed correctly and children extracted without
|
|
errors.
|
|
"""
|
|
|
|
RESOURCES = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resources", "datalake") # noqa: PTH118, PTH120
|
|
|
|
def _load_fixture_as_dataframe(self, filename):
|
|
path = os.path.join(self.RESOURCES, filename) # noqa: PTH118
|
|
with open(path, "rb") as f: # noqa: PTH123
|
|
data = json.loads(f.read())
|
|
if isinstance(data, dict):
|
|
data = [data]
|
|
return pd.DataFrame.from_records(data)
|
|
|
|
def _parsed_columns(self, filename):
|
|
df = self._load_fixture_as_dataframe(filename)
|
|
return {col.name.root: col for col in GenericDataFrameColumnParser._get_columns(df)}
|
|
|
|
def test_dict_valued_columns_typed_as_json(self):
|
|
cols = self._parsed_columns("dbt_catalog.json")
|
|
assert cols["metadata"].dataType == DataType.JSON
|
|
assert cols["nodes"].dataType == DataType.JSON
|
|
assert cols["sources"].dataType == DataType.JSON
|
|
|
|
def test_null_column_typed_as_string(self):
|
|
cols = self._parsed_columns("dbt_catalog.json")
|
|
assert cols["errors"].dataType == DataType.STRING
|
|
|
|
def test_non_empty_dict_column_has_children(self):
|
|
cols = self._parsed_columns("dbt_catalog.json")
|
|
assert cols["nodes"].children is not None and len(cols["nodes"].children) > 0
|
|
|
|
def test_empty_dict_columns_typed_as_json_not_string(self):
|
|
cols = self._parsed_columns("dbt_manifest.json")
|
|
for name in ("metrics", "groups", "disabled", "group_map", "saved_queries", "semantic_models", "unit_tests"):
|
|
assert cols[name].dataType == DataType.JSON, f"column '{name}': expected JSON, got {cols[name].dataType}"
|
|
|
|
def test_empty_dict_columns_have_no_children(self):
|
|
cols = self._parsed_columns("dbt_manifest.json")
|
|
for name in ("metrics", "groups", "disabled", "group_map", "saved_queries", "semantic_models", "unit_tests"):
|
|
children = cols[name].children
|
|
assert not children, f"column '{name}' should have no children"
|
|
|
|
|
|
class TestDbtSingleObjectJsonIngestion:
|
|
"""Single-object JSON files (e.g. dbt artifacts) are wrapped into a 1-row DataFrame
|
|
where every top-level key becomes a column with a Python dict value. The column parser
|
|
must correctly type all columns — including empty-dict columns — without errors."""
|
|
|
|
@staticmethod
|
|
def _make_catalog_df():
|
|
return pd.DataFrame(
|
|
[
|
|
{
|
|
"metadata": {"dbt_version": "1.5.0", "generated_at": "2024-01-01"},
|
|
"nodes": {"model.Project.tbl": {"name": "tbl", "description": "test"}},
|
|
"sources": {},
|
|
"errors": None,
|
|
}
|
|
]
|
|
)
|
|
|
|
@staticmethod
|
|
def _make_manifest_df():
|
|
return pd.DataFrame(
|
|
[
|
|
{
|
|
"metadata": {"dbt_version": "1.5.0"},
|
|
"nodes": {"model.Project.tbl": {"name": "tbl"}},
|
|
"sources": {},
|
|
"metrics": {},
|
|
"groups": {},
|
|
"disabled": {},
|
|
"group_map": {},
|
|
"saved_queries": {},
|
|
"semantic_models": {},
|
|
"unit_tests": {},
|
|
}
|
|
]
|
|
)
|
|
|
|
def test_catalog_column_types(self):
|
|
df = self._make_catalog_df()
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "metadata") == DataType.JSON
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "nodes") == DataType.JSON
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "sources") == DataType.JSON
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, "errors") == DataType.STRING
|
|
|
|
def test_manifest_empty_dict_columns_typed_as_json(self):
|
|
df = self._make_manifest_df()
|
|
for col in ("metrics", "groups", "disabled", "group_map", "saved_queries", "semantic_models", "unit_tests"):
|
|
assert GenericDataFrameColumnParser.fetch_col_types(df, col) == DataType.JSON, f"{col} should be JSON"
|
|
|
|
def test_catalog_nodes_children_extracted_without_error(self):
|
|
df = self._make_catalog_df()
|
|
nodes_col = df["nodes"].dropna()[:100]
|
|
children = GenericDataFrameColumnParser.get_children(nodes_col)
|
|
assert len(children) > 0
|
|
|
|
def test_catalog_sources_empty_dict_returns_no_children(self):
|
|
df = self._make_catalog_df()
|
|
sources_col = df["sources"].dropna()[:100]
|
|
assert GenericDataFrameColumnParser.get_children(sources_col) == []
|
|
|
|
|
|
class TestCSVQuotedHeaderFix(TestCase):
|
|
"""Test CSV parsing with quoted header fix for malformed CSV files"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""Set up a DSVDataFrameReader instance for testing"""
|
|
from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
|
|
LocalConfig,
|
|
)
|
|
|
|
cls.csv_reader = DSVDataFrameReader(config_source=LocalConfig(), client=None, separator=",")
|
|
cls.tsv_reader = DSVDataFrameReader(config_source=LocalConfig(), client=None, separator="\t")
|
|
|
|
def test_normal_csv_no_fix_applied(self):
|
|
"""Test that normal CSV files with proper headers are not modified"""
|
|
df = pd.DataFrame(
|
|
{
|
|
"Year": [2024, 2024, 2024],
|
|
"Industry_code": ["99999", "99999", "99999"],
|
|
"Industry_name": ["All industries", "All industries", "All industries"],
|
|
"Units": [
|
|
"Dollars (millions)",
|
|
"Dollars (millions)",
|
|
"Dollars (millions)",
|
|
],
|
|
"Value": [979594, 838626, 112188],
|
|
}
|
|
)
|
|
|
|
result = self.csv_reader._fix_malformed_quoted_chunk([df], ",")
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(len(result[0].columns), 5)
|
|
self.assertEqual(
|
|
list(result[0].columns),
|
|
["Year", "Industry_code", "Industry_name", "Units", "Value"],
|
|
)
|
|
self.assertEqual(len(result[0]), 3)
|
|
|
|
def test_malformed_csv_quoted_header_fix_applied(self):
|
|
"""Test that malformed CSV with quoted header row is properly fixed"""
|
|
malformed_header = "managementLevel,businessAllocation2Key,validFrom,fillable,businessAllocation1English"
|
|
df = pd.DataFrame(columns=[malformed_header])
|
|
|
|
result = self.csv_reader._fix_malformed_quoted_chunk([df], ",")
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(len(result[0].columns), 5)
|
|
self.assertEqual(
|
|
list(result[0].columns),
|
|
[
|
|
"managementLevel",
|
|
"businessAllocation2Key",
|
|
"validFrom",
|
|
"fillable",
|
|
"businessAllocation1English",
|
|
],
|
|
)
|
|
|
|
def test_single_column_csv_without_separator_no_fix(self):
|
|
"""Test that single column CSV without separator in name is not modified"""
|
|
df = pd.DataFrame(columns=["single_column_name"])
|
|
|
|
result = self.csv_reader._fix_malformed_quoted_chunk([df], ",")
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(list(result[0].columns), ["single_column_name"])
|
|
|
|
def test_quoted_header_with_special_characters(self):
|
|
"""Test parsing quoted header containing special characters"""
|
|
malformed_header = '"col1","col2 with spaces","col3&special","col4/slash"'
|
|
df = pd.DataFrame(columns=[malformed_header])
|
|
|
|
result = self.csv_reader._fix_malformed_quoted_chunk([df], ",")
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(len(result[0].columns), 4)
|
|
self.assertEqual(
|
|
list(result[0].columns),
|
|
["col1", "col2 with spaces", "col3&special", "col4/slash"],
|
|
)
|
|
|
|
def test_tsv_malformed_header_fix(self):
|
|
"""Test that malformed TSV with tab-separated quoted header is properly fixed"""
|
|
malformed_header = "col1\tcol2\tcol3\tcol4"
|
|
df = pd.DataFrame(columns=[malformed_header])
|
|
|
|
result = self.tsv_reader._fix_malformed_quoted_chunk([df], "\t")
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(len(result[0].columns), 4)
|
|
self.assertEqual(list(result[0].columns), ["col1", "col2", "col3", "col4"])
|
|
|
|
def test_empty_chunk_list_returns_empty(self):
|
|
"""Test that empty chunk list returns empty list"""
|
|
result = self.csv_reader._fix_malformed_quoted_chunk([], ",")
|
|
self.assertEqual(result, [])
|