llm/tests/test_plugins.py
Simon Willison ae87f978bd Moved iter_prompt from Response to Model, moved a lot of other stuff
- Moved a whole bunch of things from llm/cli.py into llm/__init__.py
- Switched plugin listings to use importlib.metadata to avoid deprecation warning
- iter_prompt() is now a method on Model, not on Response
2023-07-10 07:45:11 -07:00

40 lines
1.1 KiB
Python

from click.testing import CliRunner
import click
import importlib
import llm
from llm import cli, hookimpl, plugins
def test_register_commands():
importlib.reload(cli)
def plugin_names():
return [plugin["name"] for plugin in llm.get_plugins()]
assert "HelloWorldPlugin" not in plugin_names()
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 "HelloWorldPlugin" in plugin_names()
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 "HelloWorldPlugin" not in plugin_names()