chore: initial commit

Co-authored-by: Lifei Zhou <lifei@squareup.com>
Co-authored-by: Mic Neale <micn@tbd.email>
Co-authored-by: Lily Delalande <ldelalande@squareup.com>
Co-authored-by: Bradley Axen <baxen@squareup.com>
Co-authored-by: Andy Lane <alane@squareup.com>
Co-authored-by: Elena Zherdeva <ezherdeva@squareup.com>
Co-authored-by: Zaki Ali <zaki@squareup.com>
Co-authored-by: Salman Mohammed <smohammed@squareup.com>

tests/cli/prompt/test_goose_prompt_session.py

@@ -0,0 +1,47 @@
+from unittest.mock import patch
+
+import pytest
+from goose.cli.prompt.goose_prompt_session import GoosePromptSession
+from goose.cli.prompt.user_input import PromptAction, UserInput
+
+
+@pytest.fixture
+def mock_prompt_session():
+    with patch("prompt_toolkit.PromptSession") as mock_prompt_session:
+        yield mock_prompt_session
+
+
+def test_get_save_session_name(mock_prompt_session):
+    mock_prompt_session.prompt.return_value = "my_session"
+    goose_prompt_session = GoosePromptSession(mock_prompt_session)
+
+    assert goose_prompt_session.get_save_session_name() == "my_session"
+
+
+def test_get_user_input_to_continue(mock_prompt_session):
+    mock_prompt_session.prompt.return_value = "input_value"
+    goose_prompt_session = GoosePromptSession(mock_prompt_session)
+
+    user_input = goose_prompt_session.get_user_input()
+
+    assert user_input == UserInput(PromptAction.CONTINUE, "input_value")
+
+
+@pytest.mark.parametrize("exit_input", ["exit", ":q"])
+def test_get_user_input_to_exit(exit_input, mock_prompt_session):
+    mock_prompt_session.prompt.return_value = exit_input
+    goose_prompt_session = GoosePromptSession(mock_prompt_session)
+
+    user_input = goose_prompt_session.get_user_input()
+
+    assert user_input == UserInput(PromptAction.EXIT)
+
+
+@pytest.mark.parametrize("error", [EOFError, KeyboardInterrupt])
+def test_get_user_input_to_exit_when_error_occurs(error, mock_prompt_session):
+    mock_prompt_session.prompt.side_effect = error
+    goose_prompt_session = GoosePromptSession(mock_prompt_session)
+
+    user_input = goose_prompt_session.get_user_input()
+
+    assert user_input == UserInput(PromptAction.EXIT)

tests/cli/prompt/test_lexer.py

@@ -0,0 +1,253 @@
+from goose.cli.prompt.lexer import (
+    PromptLexer,
+    command_itself,
+    completion_for_command,
+    value_for_command,
+)
+from prompt_toolkit.document import Document
+
+
+# Helper function to create a Document and lexer instance
+def create_lexer_and_document(commands, text):
+    lexer = PromptLexer(commands)
+    document = Document(text)
+    return lexer, document
+
+
+# Test cases
+def test_lex_document_command():
+    lexer, document = create_lexer_and_document(["file"], "/file:example.txt")
+    tokens = lexer.lex_document(document)
+    expected_tokens = [("class:command", "/file:"), ("class:parameter", "example.txt")]
+    assert tokens(0) == expected_tokens
+
+
+def test_lex_document_partial_command():
+    lexer, document = create_lexer_and_document(["file"], "/fi")
+    tokens = lexer.lex_document(document)
+    expected_tokens = [("class:command", "/fi")]
+    assert tokens(0) == expected_tokens
+
+
+def test_lex_document_with_text():
+    lexer, document = create_lexer_and_document(["file"], "Some text /file:example.txt")
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:text", "S"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "e"),
+        ("class:text", " "),
+        ("class:text", "t"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+        ("class:text", " "),
+        ("class:command", "/file:"),
+        ("class:parameter", "example.txt"),
+    ]
+    assert tokens(0) == expected_tokens
+
+
+def test_lex_document_with_command_in_middle():
+    lexer, document = create_lexer_and_document(["file"], "Some text /file:example.txt more text")
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:text", "S"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "e"),
+        ("class:text", " "),
+        ("class:text", "t"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+        ("class:text", " "),
+        ("class:command", "/file:"),
+        ("class:parameter", "example.txt"),
+        ("class:text", " "),
+        ("class:text", "m"),
+        ("class:text", "o"),
+        ("class:text", "r"),
+        ("class:text", "e"),
+        ("class:text", " "),
+        ("class:text", "t"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+    ]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_multiple_commands():
+    lexer, document = create_lexer_and_document(
+        ["command", "anothercommand"],
+        "/command:example1.txt more text /anothercommand:example2.txt",
+    )
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:command", "/command:"),
+        ("class:parameter", "example1.txt"),
+        ("class:text", " "),
+        ("class:text", "m"),
+        ("class:text", "o"),
+        ("class:text", "r"),
+        ("class:text", "e"),
+        ("class:text", " "),
+        ("class:text", "t"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+        ("class:text", " "),
+        ("class:command", "/anothercommand:"),
+        ("class:parameter", "example2.txt"),
+    ]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_multiple_same_commands():
+    lexer, document = create_lexer_and_document(
+        ["command"],
+        "/command:example1.txt more text /command:example2.txt",
+    )
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:command", "/command:"),
+        ("class:parameter", "example1.txt"),
+        ("class:text", " "),
+        ("class:text", "m"),
+        ("class:text", "o"),
+        ("class:text", "r"),
+        ("class:text", "e"),
+        ("class:text", " "),
+        ("class:text", "t"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+        ("class:text", " "),
+        ("class:command", "/command:"),
+        ("class:parameter", "example2.txt"),
+    ]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_two_half_commands():
+    lexer, document = create_lexer_and_document(
+        ["command"],
+        "/comma /com",
+    )
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:text", "/"),
+        ("class:text", "c"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "m"),
+        ("class:text", "a"),
+        ("class:text", " "),
+        ("class:command", "/com"),
+    ]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_command_attached_to_pre_string():
+    lexer, document = create_lexer_and_document(
+        ["command"],
+        "some/command:example.txt",
+    )
+    expected_tokens = [
+        ("class:text", "s"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "e"),
+        ("class:text", "/"),
+        ("class:text", "c"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "m"),
+        ("class:text", "a"),
+        ("class:text", "n"),
+        ("class:text", "d"),
+        ("class:text", ":"),
+        ("class:text", "e"),
+        ("class:text", "x"),
+        ("class:text", "a"),
+        ("class:text", "m"),
+        ("class:text", "p"),
+        ("class:text", "l"),
+        ("class:text", "e"),
+        ("class:text", "."),
+        ("class:text", "t"),
+        ("class:text", "x"),
+        ("class:text", "t"),
+    ]
+    tokens = lexer.lex_document(document)
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_partial_command_attached_to_pre_string():
+    lexer, document = create_lexer_and_document(
+        ["command"],
+        "some/com",
+    )
+    tokens = lexer.lex_document(document)
+    expected_tokens = [
+        ("class:text", "s"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+        ("class:text", "e"),
+        ("class:text", "/"),
+        ("class:text", "c"),
+        ("class:text", "o"),
+        ("class:text", "m"),
+    ]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_no_command():
+    lexer, document = create_lexer_and_document([], "Some random text")
+    tokens = lexer.lex_document(document)
+    expected_tokens = [("class:text", character) for character in "Some random text"]
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_lex_document_ending_char_of_parameter_is_symbol():
+    lexer, document = create_lexer_and_document(
+        ["command"],
+        "/command:example.txt/",
+    )
+    expected_tokens = [
+        ("class:command", "/command:"),
+        ("class:parameter", "example.txt/"),
+    ]
+    tokens = lexer.lex_document(document)
+    actual_tokens = list(tokens(0))
+    assert actual_tokens == expected_tokens
+
+
+def test_command_itself():
+    pattern = command_itself("file:")
+    matches = pattern.match("/file:example.txt")
+    assert matches is not None
+    assert matches.group(1) == "/file:"
+
+
+def test_value_for_command():
+    pattern = value_for_command("file:")
+    matches = pattern.search("/file:example.txt")
+    assert matches is not None
+    assert matches.group(1) == "example.txt"
+
+
+def test_completion_for_command():
+    pattern = completion_for_command("file:")
+    matches = pattern.search("/file:")
+    assert matches is not None
+    assert matches.group(1) == "file:"

tests/cli/prompt/test_prompt_validator.py

@@ -0,0 +1,37 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from goose.cli.prompt.prompt_validator import PromptValidator
+from prompt_toolkit.validation import ValidationError
+
+
+@pytest.fixture
+def validator():
+    return PromptValidator()
+
+
+@patch("prompt_toolkit.document.Document.text")
+def test_validate_should_not_raise_error_when_input_is_none(document, validator):
+    try:
+        validator.validate(create_mock_document(None))
+    except Exception as e:
+        pytest.fail(f"An error was raised: {e}")
+
+
+@patch("prompt_toolkit.document.Document.text", return_value="user typed something")
+def test_validate_should_not_raise_error_when_user_has_input(document, validator):
+    try:
+        validator.validate(create_mock_document("user typed something"))
+    except Exception as e:
+        pytest.fail(f"An error was raised: {e}")
+
+
+def test_validate_should_raise_validation_error_when_user_has_empty_input(validator):
+    with pytest.raises(ValidationError):
+        validator.validate(create_mock_document(""))
+
+
+def create_mock_document(text: str) -> MagicMock:
+    document = MagicMock()
+    document.text = text
+    return document

tests/cli/prompt/test_user_input.py

@@ -0,0 +1,15 @@
+from goose.cli.prompt.user_input import PromptAction, UserInput
+
+
+def test_user_input_with_action_continue():
+    input = UserInput(action=PromptAction.CONTINUE, text="Hello")
+    assert input.to_continue() is True
+    assert input.to_exit() is False
+    assert input.text == "Hello"
+
+
+def test_user_input_with_action_exit():
+    input = UserInput(action=PromptAction.EXIT)
+    assert input.to_continue() is False
+    assert input.to_exit() is True
+    assert input.text is None