sqlite-utils/tests/test_plugins.py
Alex Garcia 3f80a02698
prepare_connection plugin hook
Closes:
- #574

Refs #567

---------

Co-authored-by: Simon Willison <swillison@gmail.com>
2023-07-22 15:59:08 -07:00

79 lines
2.2 KiB
Python

from click.testing import CliRunner
import click
import importlib
from sqlite_utils import cli, Database, hookimpl, plugins
def test_register_commands():
importlib.reload(cli)
assert plugins.get_plugins() == []
class HelloWorldPlugin:
__name__ = "HelloWorldPlugin"
@hookimpl
def register_commands(self, cli):
@cli.command(name="hello-world")
def hello_world():
"Print hello world"
click.echo("Hello world!")
try:
plugins.pm.register(HelloWorldPlugin(), name="HelloWorldPlugin")
importlib.reload(cli)
assert plugins.get_plugins() == [
{"name": "HelloWorldPlugin", "hooks": ["register_commands"]}
]
runner = CliRunner()
result = runner.invoke(cli.cli, ["hello-world"])
assert result.exit_code == 0
assert result.output == "Hello world!\n"
finally:
plugins.pm.unregister(name="HelloWorldPlugin")
importlib.reload(cli)
assert plugins.get_plugins() == []
def test_prepare_connection():
importlib.reload(cli)
assert plugins.get_plugins() == []
class HelloFunctionPlugin:
__name__ = "HelloFunctionPlugin"
@hookimpl
def prepare_connection(self, conn):
conn.create_function("hello", 1, lambda name: f"Hello, {name}!")
db = Database(memory=True)
functions = db.execute(
"select distinct name from pragma_function_list order by 1"
).fetchall()
assert "hello" not in functions
try:
plugins.pm.register(HelloFunctionPlugin(), name="HelloFunctionPlugin")
assert plugins.get_plugins() == [
{"name": "HelloFunctionPlugin", "hooks": ["prepare_connection"]}
]
db = Database(memory=True)
functions = [
row[0]
for row in db.execute(
"select distinct name from pragma_function_list order by 1"
).fetchall()
]
assert "hello" in functions
result = db.execute('select hello("world")').fetchone()[0]
assert result == "Hello, world!"
finally:
plugins.pm.unregister(name="HelloFunctionPlugin")
assert plugins.get_plugins() == []