def test_env_init():
env = Environment(
f"_test_{rand_str()}",
metadata_storage="sqlite://",
settings=SnapflowSettings(add_core_module=False),
)
env_init(env)
def get_env(key="_test", db_url=None):
if db_url is None:
db_url = get_tmp_sqlite_db_url()
env = Environment(key=key, metadata_storage=db_url)
env.add_module(core)
env.add_schema(Customer)
env.add_schema(Metric)
return env
def test_example():
env = Environment(metadata_storage="sqlite://")
g = Graph(env)
env.add_module(core)
df = pd.DataFrame({"a": range(10), "b": range(10)})
g.create_node(key="n1", pipe="extract_dataframe", config={"dataframe": df})
output = env.produce("n1", g)
assert_almost_equal(output.as_dataframe(), df)
def test_simple_import():
dburl = get_tmp_sqlite_db_url()
env = Environment(metadata_storage=dburl)
g = Graph(env)
env.add_module(core)
df = pd.DataFrame({"a": range(10), "b": range(10)})
g.create_node(key="n1", function="import_dataframe", params={"dataframe": df})
blocks = env.produce("n1", g)
assert_almost_equal(blocks[0].as_dataframe(), df, check_dtype=False)
def app(ctx, debug: bool = False, metadata: Optional[str] = None):
"""Modern Data Pipelines"""
logger.warning("The snapflow CLI is experimental and not officially supported yet")
if debug:
logger.add(sys.stderr, level="DEBUG")
else:
logger.add(sys.stderr, level="INFO")
env = current_env()
if env is None:
env = Environment(metadata_storage=metadata)
logger.info(f"Using environment '{env.metadata_storage.url}'")
ctx.obj = env
def make_test_env(**kwargs) -> Environment:
if "metadata_storage" not in kwargs:
url = "sqlite://"
metadata_storage = Storage.from_url(url)
kwargs["metadata_storage"] = metadata_storage
env = Environment(**kwargs)
test_module = SnapflowModule(
"_test",
schemas=[TestSchema1, TestSchema2, TestSchema3, TestSchema4],
)
env.add_module(test_module)
return env
def test_default_module():
DEFAULT_LOCAL_MODULE.library.snaps = {}
@Snap
def s1():
pass
assert len(DEFAULT_LOCAL_MODULE.library.snaps) == 1
assert DEFAULT_LOCAL_MODULE.get_snap("s1") is s1
env = Environment()
env.add_snap(s1)
assert env.get_snap("s1") is s1
def make_test_env(**kwargs) -> Environment:
if "metadata_storage" not in kwargs:
url = get_tmp_sqlite_db_url()
metadata_storage = Storage.from_url(url)
kwargs["metadata_storage"] = metadata_storage
env = Environment(settings=SnapflowSettings(abort_on_function_error=True), **kwargs)
test_module = SnapflowModule(
"_test",
)
for schema in [TestSchema1, TestSchema2, TestSchema3, TestSchema4]:
env.add_schema(schema)
env.add_module(test_module)
return env
def test_default_module():
DEFAULT_LOCAL_MODULE.library.functions = {}
@datafunction
def s1():
pass
assert len(DEFAULT_LOCAL_MODULE.library.functions) == 1
assert DEFAULT_LOCAL_MODULE.get_function("s1") is s1
env = Environment()
env.add_function(s1)
assert env.get_function("s1") is s1
def test_env_init():
from . import _test_module
# Test module / components
env = Environment("_test",
metadata_storage="sqlite://",
initial_modules=[])
with env.session_scope() as sess:
assert len(env.get_module_order()) == 1
env.add_module(_test_module)
assert env.get_module_order() == [
env.get_local_module().name,
_test_module.name,
]
assert env.get_schema("TestSchema",
sess) is _test_module.schemas.TestSchema
assert env.get_pipe("test_sql") is _test_module.pipes.test_sql
# Test runtime / storage
env.add_storage("postgresql://test")
assert len(env.storages) == 2 # added plus default local memory
assert len(env.runtimes) == 3 # added plus default local python # TODO
def get_env():
env = Environment(metadata_storage="sqlite://")
env.add_module(core)
env.add_schema(Customer)
env.add_schema(Metric)
return env