Merging upstream version 1.14.2.

Signed-off-by: Daniel Baumann <daniel@debian.org>

tests/test_completion_engine.py

@@ -357,6 +357,18 @@ def test_sub_select_multiple_col_name_completion():
     )
 
 
+def test_suggested_multiple_column_names():
+    suggestions = suggest_type("SELECT id, from users", "SELECT id, ")
+    assert sorted_dicts(suggestions) == sorted_dicts(
+        [
+            {"type": "column", "tables": [(None, "users", None)]},
+            {"type": "function", "schema": []},
+            {"type": "alias", "aliases": ["users"]},
+            {"type": "keyword"},
+        ]
+    )
+
+
 def test_sub_select_dot_col_name_completion():
     suggestions = suggest_type("SELECT * FROM (SELECT t. FROM tabl t", "SELECT * FROM (SELECT t.")
     assert sorted_dicts(suggestions) == sorted_dicts(

tests/test_llm_special.py

@@ -0,0 +1,162 @@
+import pytest
+from unittest.mock import patch
+from litecli.packages.special.llm import handle_llm, FinishIteration, USAGE
+
+
+@patch("litecli.packages.special.llm.initialize_llm")
+@patch("litecli.packages.special.llm.llm", new=None)
+def test_llm_command_without_install(mock_initialize_llm, executor):
+    """
+    Test that handle_llm initializes llm when it is None and raises FinishIteration.
+    """
+    test_text = r"\llm"
+    cur_mock = executor
+
+    with pytest.raises(FinishIteration) as exc_info:
+        handle_llm(test_text, cur_mock)
+
+    mock_initialize_llm.assert_called_once()
+    assert exc_info.value.args[0] is None
+
+
+@patch("litecli.packages.special.llm.llm")
+def test_llm_command_without_args(mock_llm, executor):
+    r"""
+    Invoking \llm without any arguments should print the usage and raise
+    FinishIteration.
+    """
+    assert mock_llm is not None
+    test_text = r"\llm"
+    cur_mock = executor
+
+    with pytest.raises(FinishIteration) as exc_info:
+        handle_llm(test_text, cur_mock)
+
+    assert exc_info.value.args[0] == [(None, None, None, USAGE)]
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.run_external_cmd")
+def test_llm_command_with_c_flag(mock_run_cmd, mock_llm, executor):
+    # Suppose the LLM returns some text without fenced SQL
+    mock_run_cmd.return_value = (0, "Hello, I have no SQL for you today.")
+
+    test_text = r"\llm -c 'Something interesting?'"
+
+    with pytest.raises(FinishIteration) as exc_info:
+        handle_llm(test_text, executor)
+
+    # We expect no code fence => FinishIteration with that output
+    assert exc_info.value.args[0] == [(None, None, None, "Hello, I have no SQL for you today.")]
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.run_external_cmd")
+def test_llm_command_with_c_flag_and_fenced_sql(mock_run_cmd, mock_llm, executor):
+    # The luscious SQL is inside triple backticks
+    return_text = "Here is your query:\n" "```sql\nSELECT * FROM table;\n```"
+    mock_run_cmd.return_value = (0, return_text)
+
+    test_text = r"\llm -c 'Rewrite the SQL without CTE'"
+
+    result, sql = handle_llm(test_text, executor)
+
+    # We expect the function to return (result, sql), but result might be "" if verbose is not set
+    # By default, `verbose` is false unless text has something like \llm --verbose?
+    # The function code: return result if verbose else "", sql
+    # Our test_text doesn't set verbose => we expect "" for the returned context.
+    assert result == ""
+    assert sql == "SELECT * FROM table;"
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.run_external_cmd")
+def test_llm_command_known_subcommand(mock_run_cmd, mock_llm, executor):
+    """
+    If the parts[0] is in LLM_CLI_COMMANDS, we do NOT capture output, we just call run_external_cmd
+    and then raise FinishIteration.
+    """
+    # Let's assume 'models' is in LLM_CLI_COMMANDS
+    test_text = r"\llm models"
+
+    with pytest.raises(FinishIteration) as exc_info:
+        handle_llm(test_text, executor)
+
+    # We check that run_external_cmd was called with these arguments:
+    mock_run_cmd.assert_called_once_with("llm", "models", restart_cli=False)
+    # And the function should raise FinishIteration(None)
+    assert exc_info.value.args[0] is None
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.run_external_cmd")
+def test_llm_command_with_install_flag(mock_run_cmd, mock_llm, executor):
+    """
+    If 'install' or 'uninstall' is in the parts, we do not capture output but restart the CLI.
+    """
+    test_text = r"\llm install openai"
+
+    with pytest.raises(FinishIteration) as exc_info:
+        handle_llm(test_text, executor)
+
+    # We expect a restart
+    mock_run_cmd.assert_called_once_with("llm", "install", "openai", restart_cli=True)
+    assert exc_info.value.args[0] is None
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.ensure_litecli_template")
+@patch("litecli.packages.special.llm.sql_using_llm")
+def test_llm_command_with_prompt(mock_sql_using_llm, mock_ensure_template, mock_llm, executor):
+    r"""
+    \llm prompt "some question"
+    Should use context, capture output, and call sql_using_llm.
+    """
+    # Mock out the return from sql_using_llm
+    mock_sql_using_llm.return_value = ("context from LLM", "SELECT 1;")
+
+    test_text = r"\llm prompt 'Magic happening here?'"
+    context, sql = handle_llm(test_text, executor)
+
+    # ensure_litecli_template should be called
+    mock_ensure_template.assert_called_once()
+    # sql_using_llm should be called with question=arg, which is "prompt 'Magic happening here?'"
+    # Actually, the question is the entire "prompt 'Magic happening here?'" minus the \llm
+    # But in the function we do parse shlex.split.
+    mock_sql_using_llm.assert_called()
+    assert context == ""
+    assert sql == "SELECT 1;"
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.ensure_litecli_template")
+@patch("litecli.packages.special.llm.sql_using_llm")
+def test_llm_command_question_with_context(mock_sql_using_llm, mock_ensure_template, mock_llm, executor):
+    """
+    If arg doesn't contain any known command, it's treated as a question => capture output + context.
+    """
+    mock_sql_using_llm.return_value = ("You have context!", "SELECT 2;")
+
+    test_text = r"\llm 'Top 10 downloads by size.'"
+    context, sql = handle_llm(test_text, executor)
+
+    mock_ensure_template.assert_called_once()
+    mock_sql_using_llm.assert_called()
+    assert context == ""
+    assert sql == "SELECT 2;"
+
+
+@patch("litecli.packages.special.llm.llm")
+@patch("litecli.packages.special.llm.ensure_litecli_template")
+@patch("litecli.packages.special.llm.sql_using_llm")
+def test_llm_command_question_verbose(mock_sql_using_llm, mock_ensure_template, mock_llm, executor):
+    r"""
+    Invoking \llm+ returns the context and the SQL query.
+    """
+    mock_sql_using_llm.return_value = ("Verbose context, oh yeah!", "SELECT 42;")
+
+    test_text = r"\llm+ 'Top 10 downloads by size.'"
+    context, sql = handle_llm(test_text, executor)
+
+    assert context == "Verbose context, oh yeah!"
+    assert sql == "SELECT 42;"

tests/test_main.py

@@ -2,6 +2,8 @@ import os
 from collections import namedtuple
 from textwrap import dedent
 import shutil
+from datetime import datetime
+from unittest.mock import patch
 
 import click
 from click.testing import CliRunner
@@ -267,3 +269,64 @@ def test_startup_commands(executor):
     ]
 
     # implement tests on executions of the startupcommands
+
+
+@patch("litecli.main.datetime")  # Adjust if your module path is different
+def test_get_prompt(mock_datetime):
+    # We'll freeze time at 2025-01-20 13:37:42 for comedic effect.
+    # Because "leet" times call for 13:37!
+    frozen_time = datetime(2025, 1, 20, 13, 37, 42)
+    mock_datetime.now.return_value = frozen_time
+    # Ensure `datetime` class is still accessible for strftime usage
+    mock_datetime.datetime = datetime
+
+    # Instantiate and connect
+    lc = LiteCli()
+    lc.connect("/tmp/litecli_test.db")
+
+    # 1. Test \d => full path to the DB
+    assert lc.get_prompt(r"\d") == "/tmp/litecli_test.db"
+
+    # 2. Test \f => basename of the DB
+    #    (because "f" stands for "filename", presumably!)
+    assert lc.get_prompt(r"\f") == "litecli_test.db"
+
+    # 3. Test \_ => single space
+    assert lc.get_prompt(r"Hello\_World") == "Hello World"
+
+    # 4. Test \n => newline
+    #    Just to be sure we're only inserting a newline,
+    #    we can check length or assert the presence of "\n".
+    expected = f"Line1{os.linesep}Line2"
+    assert lc.get_prompt(r"Line1\nLine2") == expected
+
+    # 5. Test date/time placeholders (with frozen time):
+    #    \D => e.g. 'Mon Jan 20 13:37:42 2025'
+    expected_date_str = frozen_time.strftime("%a %b %d %H:%M:%S %Y")
+    assert lc.get_prompt(r"\D") == expected_date_str
+
+    # 6. Test \m => minutes
+    assert lc.get_prompt(r"\m") == "37"
+
+    # 7. Test \P => AM/PM
+    #    13:37 is PM
+    assert lc.get_prompt(r"\P") == "PM"
+
+    # 8. Test \R => 24-hour format hour
+    assert lc.get_prompt(r"\R") == "13"
+
+    # 9. Test \r => 12-hour format hour
+    #    13:37 is 01 in 12-hour format
+    assert lc.get_prompt(r"\r") == "01"
+
+    # 10. Test \s => seconds
+    assert lc.get_prompt(r"\s") == "42"
+
+    # 11. Test when dbname is None => (none)
+    lc.connect(None)  # Simulate no DB connection
+    assert lc.get_prompt(r"\d") == "(none)"
+    assert lc.get_prompt(r"\f") == "(none)"
+
+    # 12. Windows path
+    lc.connect("C:\\Users\\litecli\\litecli_test.db")
+    assert lc.get_prompt(r"\d") == "C:\\Users\\litecli\\litecli_test.db"

tests/test_sqlexecute.py

@@ -38,13 +38,15 @@ def test_binary(executor):
 
 
 ## Failing in Travis for some unknown reason.
-# @dbtest
-# def test_table_and_columns_query(executor):
-#     run(executor, "create table a(x text, y text)")
-#     run(executor, "create table b(z text)")
+@dbtest
+def test_table_and_columns_query(executor):
+    run(executor, "create table a(x text, y text)")
+    run(executor, "create table b(z text)")
+    run(executor, "create table t(t text)")
 
-#     assert set(executor.tables()) == set([("a",), ("b",)])
-#     assert set(executor.table_columns()) == set([("a", "x"), ("a", "y"), ("b", "z")])
+    assert set(executor.tables()) == set([("a",), ("b",), ("t",)])
+    assert set(executor.table_columns()) == set([("a", "x"), ("a", "y"), ("b", "z"), ("t", "t")])
+    assert set(executor.table_columns()) == set([("a", "x"), ("a", "y"), ("b", "z"), ("t", "t")])
 
 
 @dbtest