Skip to content

Commit

Permalink
Move tests to pipeline extra
Browse files Browse the repository at this point in the history
  • Loading branch information
Sultan Iman committed Feb 7, 2024
1 parent 91e5553 commit ebe6ced
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 161 deletions.
161 changes: 0 additions & 161 deletions tests/load/pipeline/test_pydantic_models.py

This file was deleted.

154 changes: 154 additions & 0 deletions tests/pipeline/test_pipeline_extra.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,3 +200,157 @@ def generic(start=8):

pipeline = dlt.pipeline(destination="duckdb")
pipeline.run(generic(), loader_file_format=file_format)


class Child(BaseModel):
child_attribute: str
optional_child_attribute: Optional[str] = None


@pytest.mark.parametrize(
"destination_config",
destinations_configs(default_sql_configs=True, subset=["duckdb"]),
ids=lambda x: x.name,
)
def test_flattens_model_when_skip_complex_types_is_set(
destination_config: DestinationTestConfiguration,
) -> None:
class Parent(BaseModel):
child: Child
optional_parent_attribute: Optional[str] = None
dlt_config: ClassVar[DltConfig] = {"skip_complex_types": True}

example_data = {
"optional_parent_attribute": None,
"child": {
"child_attribute": "any string",
"optional_child_attribute": None,
},
}

@dlt.resource
def res():
yield [example_data]

@dlt.source(max_table_nesting=1)
def src():
yield res()

p = destination_config.setup_pipeline("example", full_refresh=True)
p.run(src(), table_name="items", columns=Parent)

with p.sql_client() as client:
with client.execute_query("SELECT * FROM items") as cursor:
loaded_values = {
col[0]: val
for val, col in zip(cursor.fetchall()[0], cursor.description)
if col[0] not in ("_dlt_id", "_dlt_load_id")
}
assert loaded_values == {
"child__child_attribute": "any string",
"child__optional_child_attribute": None,
"optional_parent_attribute": None,
}

keys = p.default_schema.tables["items"]["columns"].keys()
columns = p.default_schema.tables["items"]["columns"]

assert keys == {
"child__child_attribute",
"child__optional_child_attribute",
"optional_parent_attribute",
"_dlt_load_id",
"_dlt_id",
}

assert columns["child__child_attribute"] == {
"name": "child__child_attribute",
"data_type": "text",
"nullable": False,
}

assert columns["child__optional_child_attribute"] == {
"name": "child__optional_child_attribute",
"data_type": "text",
"nullable": True,
}

assert columns["optional_parent_attribute"] == {
"name": "optional_parent_attribute",
"data_type": "text",
"nullable": True,
}


@pytest.mark.parametrize(
"destination_config",
destinations_configs(default_sql_configs=True, subset=["duckdb"]),
ids=lambda x: x.name,
)
def test_flattens_model_when_skip_complex_types_is_not_set(
destination_config: DestinationTestConfiguration,
):
class Parent(BaseModel):
child: Child
optional_parent_attribute: Optional[str] = None
data_dictionary: Dict[str, Any] = None
dlt_config: ClassVar[DltConfig] = {"skip_complex_types": False}

example_data = {
"optional_parent_attribute": None,
"data_dictionary": {
"child_attribute": "any string",
},
"child": {
"child_attribute": "any string",
"optional_child_attribute": None,
},
}

@dlt.resource
def res():
yield [example_data]

@dlt.source(max_table_nesting=1)
def src():
yield res()

p = destination_config.setup_pipeline("example", full_refresh=True)
p.run(src(), table_name="items", columns=Parent)

with p.sql_client() as client:
with client.execute_query("SELECT * FROM items") as cursor:
loaded_values = {
col[0]: val
for val, col in zip(cursor.fetchall()[0], cursor.description)
if col[0] not in ("_dlt_id", "_dlt_load_id")
}

assert loaded_values == {
"child": '{"child_attribute":"any string","optional_child_attribute":null}',
"optional_parent_attribute": None,
"data_dictionary": '{"child_attribute":"any string"}',
}

keys = p.default_schema.tables["items"]["columns"].keys()
assert keys == {
"child",
"optional_parent_attribute",
"data_dictionary",
"_dlt_load_id",
"_dlt_id",
}

columns = p.default_schema.tables["items"]["columns"]

assert columns["optional_parent_attribute"] == {
"name": "optional_parent_attribute",
"data_type": "text",
"nullable": True,
}

assert columns["data_dictionary"] == {
"name": "data_dictionary",
"data_type": "complex",
"nullable": False,
}

0 comments on commit ebe6ced

Please sign in to comment.