Spaces:
Running
Running
File size: 4,493 Bytes
5f923cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | # Copyright 2026 The ODML Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Unit tests for the main litert-lm CLI."""
import unittest.mock
from absl.testing import absltest
from click.testing import CliRunner
from prompt_toolkit import key_binding
from litert_lm_cli import main
from litert_lm_cli import model
class MainTest(absltest.TestCase):
def test_help_shorthand(self):
runner = CliRunner()
result_help = runner.invoke(main.cli, ["--help"])
result_h = runner.invoke(main.cli, ["-h"])
self.assertEqual(result_help.exit_code, 0)
self.assertEqual(result_h.exit_code, 0)
self.assertEqual(result_help.output, result_h.output)
def test_subcommand_help_shorthand(self):
runner = CliRunner()
result_help = runner.invoke(main.cli, ["list", "--help"])
result_h = runner.invoke(main.cli, ["list", "-h"])
self.assertEqual(result_help.exit_code, 0)
self.assertEqual(result_h.exit_code, 0)
self.assertEqual(result_help.output, result_h.output)
@unittest.mock.patch(
"litert_lm_cli.model.Model.from_model_reference"
)
def test_run_with_piped_input(self, mock_from_model_ref):
mock_model = unittest.mock.MagicMock()
mock_from_model_ref.return_value = mock_model
mock_model.exists.return_value = True
runner = CliRunner()
# Mocking stdin by providing input to the runner
result = runner.invoke(
main.cli, ["run", "my-model"], input="Hello from pipe\n"
)
self.assertEqual(result.exit_code, 0)
mock_model.run_interactive.assert_called_once()
kwargs = mock_model.run_interactive.call_args.kwargs
self.assertEqual(kwargs["prompt"], "Hello from pipe")
@unittest.mock.patch(
"litert_lm_cli.model.Model.from_model_reference"
)
def test_run_with_prompt_and_piped_input(self, mock_from_model_ref):
mock_model = unittest.mock.MagicMock()
mock_from_model_ref.return_value = mock_model
mock_model.exists.return_value = True
runner = CliRunner()
# Mocking stdin by providing input to the runner
result = runner.invoke(
main.cli,
["run", "my-model", "--prompt", "Prompt arg"],
input="Hello from pipe\n",
)
self.assertEqual(result.exit_code, 0)
mock_model.run_interactive.assert_called_once()
kwargs = mock_model.run_interactive.call_args.kwargs
self.assertEqual(kwargs["prompt"], "Prompt arg\n\nHello from pipe")
@unittest.mock.patch(
"litert_lm_cli.model.Model.from_model_reference"
)
def test_run_non_tty_no_input(self, mock_from_model_ref):
mock_model = unittest.mock.MagicMock()
mock_from_model_ref.return_value = mock_model
mock_model.exists.return_value = True
runner = CliRunner()
# No input provided, isatty will be False in CliRunner
result = runner.invoke(main.cli, ["run", "my-model"])
self.assertEqual(result.exit_code, 0)
# Should return early and not start the interactive session
mock_model.run_interactive.assert_not_called()
def test_create_keybindings(self):
m = model.Model(model_id="test_model", model_path="test_path")
kb = m._create_keybindings()
self.assertIsInstance(kb, key_binding.KeyBindings)
# Check if expected keys are added.
keys = [str(b.keys) for b in kb.bindings]
# Check if enter (ControlM), c-j (ControlJ), esc+enter, c-c (ControlC).
self.assertTrue(any("ControlM" in k and "Escape" not in k for k in keys))
self.assertTrue(any("ControlJ" in k for k in keys))
self.assertTrue(any("Escape" in k and "ControlM" in k for k in keys))
self.assertTrue(any("ControlC" in k for k in keys))
def test_run_no_template_flag(self):
runner = CliRunner()
# Test that --no-template is a valid option for the run command.
# We use --help to avoid actually running the model.
result = runner.invoke(main.cli, ["run", "--help"])
self.assertEqual(result.exit_code, 0)
self.assertIn("--no-template", result.output)
if __name__ == "__main__":
absltest.main()
|