def tmp_session() -> Iterator[Session]:
"""Creates temporary Session instance for database interaction.
The database is a temporary sqlite instance created in memory.
Yields:
session: temp Session instance
"""
engine = sqlalchemy.create_engine("sqlite:///:memory:")
config = Config(config_dir=MagicMock())
config.init_db(engine=engine)
with session_scope() as session:
yield session
def _tmp_config(settings: str = "") -> Config:
config_dir = tmp_path_factory.mktemp("config")
if settings:
settings_path = config_dir / "config.toml"
settings_path.write_text(textwrap.dedent(settings))
return Config(config_dir=config_dir, settings_filename="config.toml")
def run_migrations_online():
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine and associate a connection
with the context.
"""
connectable = config.attributes.get("connection", None)
if not connectable:
# only create Engine if we don't have a Connection from the outside
moe_config = Config()
moe_config.init_db(create_tables=False)
connectable = moe_config.engine
# When connectable is already a Connection object, calling
# connect() gives us a *branched connection*.
with connectable.connect() as connection:
context.configure(connection=connection,
target_metadata=target_metadata)
with context.begin_transaction():
context.run_migrations()
def _parse_args(args: List[str], config: Config):
"""Parses the commandline arguments.
Args:
args: Arguments to parse. Should not include 'moe'.
config: User configuration for moe.
Raises:
SystemExit: No sub-commands given.
Does not include root commands such as `--version` or `--help`.
"""
moe_parser = _create_arg_parser()
# load all sub-commands
cmd_parsers = moe_parser.add_subparsers(help="command to run",
dest="command")
config.plugin_manager.hook.add_command(cmd_parsers=cmd_parsers)
parsed_args = moe_parser.parse_args(args)
# no sub-command given
if not parsed_args.command:
moe_parser.print_help(sys.stderr)
raise SystemExit(1)
_set_root_log_lvl(parsed_args)
# call the sub-command's handler within a single session
config.init_db()
with session_scope() as session:
sqlalchemy.event.listen(
session, "before_flush",
functools.partial(_edit_new_items, config=config))
sqlalchemy.event.listen(
session, "after_flush",
functools.partial(_process_new_items, config=config))
parsed_args.func(config, session, args=parsed_args)
def test_config_dir_dne(self, tmp_path):
"""Should create the config directory if it doesn't exist."""
config = Config(tmp_path / "doesn't exist")
assert config.config_dir.is_dir()
def test_config_dir_env(self, tmp_path):
"""The configuration directory can be set with an env var."""
with patch.dict("os.environ", {"MOE_CONFIG_DIR": str(tmp_path)}):
config = Config()
assert config.config_dir == tmp_path
def test_config_file_dne(self, tmp_path):
"""Should create an empty config file if it doesn't exist."""
Config(config_dir=tmp_path, settings_filename="config.toml")
assert (tmp_path / "config.toml").is_file()
def main(args: List[str] = sys.argv[1:], config: Config = None):
"""Runs the CLI."""
if not config:
config = Config()
_parse_args(args, config)