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 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 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 env_init(env: Environment):
from . import _test_module
# Test module / components
with env.md_api.begin():
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") is _test_module.schemas.TestSchema
assert env.get_snap("test_sql") is _test_module.snaps.test_sql
# Test runtime / storage
env.add_storage("postgresql://test")
assert len(env.storages) == 2 # added plus default local memory
assert len(env.runtimes) == 2 # added plus default local python
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