plesna/tests/dataplatform/test_dataplateform.py

138 lines
4.0 KiB
Python

import shutil
from pathlib import Path
import pytest
from plesna.dataplatform import DataPlateform
from plesna.models.flux import Flux
from plesna.models.transformation import Transformation
from plesna.storage.repository.fs_repository import FSRepository
FIXTURE_DIR = Path(__file__).parent.parent / Path("raw_datas")
@pytest.fixture
def repository(tmp_path) -> FSRepository:
example_src = FIXTURE_DIR
assert example_src.exists()
raw_path = Path(tmp_path) / "raw"
shutil.copytree(src=example_src.absolute(), dst=raw_path.absolute())
bronze_path = Path(tmp_path) / "bronze"
bronze_path.mkdir()
silver_path = Path(tmp_path) / "silver"
silver_path.mkdir()
return FSRepository("test", tmp_path, "test")
def test_add_repository(
repository: FSRepository,
):
dp = DataPlateform()
dp.add_repository("test", repository)
assert dp.repositories == ["test"]
assert dp.repository("test") == repository
@pytest.fixture
def foo_flux(repository: FSRepository) -> Flux:
src = {"username": repository.table("test-raw-username")}
targets = {"username": repository.table("test-bronze-username")}
def foo(sources, targets):
return {"who": "foo"}
extra_kwrds = {}
flux = Flux(
sources=src,
targets=targets,
transformation=Transformation(function=foo, extra_kwrds=extra_kwrds),
)
return flux
@pytest.fixture
def copy_flux(repository: FSRepository) -> Flux:
raw_username = {"username": repository.table("test-raw-username")}
bronze_username = {"username": repository.table("test-bronze-username")}
def copy(sources, targets):
src_path = Path(sources["username"].datas[0])
tgt_path = Path(targets["username"].datas[0])
shutil.copy(src_path, tgt_path)
return {"src_size": src_path.stat().st_size, "tgt_size": tgt_path.stat().st_size}
extra_kwrds = {}
raw_brz_copy_username = Flux(
sources=raw_username,
targets=bronze_username,
transformation=Transformation(function=copy, extra_kwrds=extra_kwrds),
)
return raw_brz_copy_username
def test_add_flux(repository: FSRepository, copy_flux: Flux):
dataplatform = DataPlateform()
dataplatform.add_repository("test", repository)
dataplatform.add_flux(name="copy_flux", flux=copy_flux)
assert dataplatform.fluxes == ["copy_flux"]
dataplatform.add_flux(name="copy_flux_bis", flux=copy_flux)
assert dataplatform.fluxes == ["copy_flux", "copy_flux_bis"]
assert dataplatform.flux("copy_flux") == copy_flux
assert dataplatform.flux("copy_flux_bis") == copy_flux
@pytest.fixture
def dataplatform(
repository: FSRepository,
foo_flux: Flux,
copy_flux: Flux,
) -> DataPlateform:
dp = DataPlateform()
dp.add_repository("test", repository)
dp.add_flux("foo", foo_flux)
dp.add_flux("raw_brz_copy_username", copy_flux)
return dp
def test_listing_content(dataplatform: DataPlateform):
assert dataplatform.repository("test").schemas() == ["test-raw", "test-bronze", "test-silver"]
assert dataplatform.repository("test").schema("test-raw").tables == [
"test-raw-username",
"test-raw-recovery",
"test-raw-salary",
]
assert dataplatform.repository("test").table("test-raw-username").partitions == ["username.csv"]
assert dataplatform.repository("test").table("test-raw-recovery").partitions == [
"2022.csv",
"2023.csv",
"2024.csv",
]
def test_content_from_graph(dataplatform: DataPlateform):
# assert dataplatform.graphset.model_dump() == {}
pass
def test_execute_flux(dataplatform: DataPlateform):
meta = dataplatform.execute_flux("foo")
assert meta.data == {"who": "foo"}
assert dataplatform.repository("test").schema("test-bronze").tables == []
meta = dataplatform.execute_flux("raw_brz_copy_username")
assert meta.data == {"src_size": 283, "tgt_size": 283}
assert dataplatform.repository("test").schema("test-bronze").tables == ["test-bronze-username"]