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>
2024-08-23 16:39:04 -07:00
commit dd126afa6c
68 changed files with 4498 additions and 0 deletions
@@ -0,0 +1,136 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from exchange import Exchange
+from goose.utils.ask import ask_an_ai, clear_exchange, replace_prompt
+
+
+# tests for `ask_an_ai`
+def test_ask_an_ai_empty_input():
+    """Test that function raises TypeError if input is empty."""
+    exchange = MagicMock(spec=Exchange)
+    with pytest.raises(TypeError):
+        ask_an_ai("", exchange)
+
+
+def test_ask_an_ai_no_history():
+    """Test the no_history functionality."""
+    exchange = MagicMock(spec=Exchange)
+    with patch("goose.utils.ask.clear_exchange") as mock_clear:
+        ask_an_ai("Test input", exchange, no_history=True)
+        mock_clear.assert_called_once_with(exchange)
+
+
+def test_ask_an_ai_prompt_replacement():
+    """Test that the prompt is replaced if provided."""
+    exchange = MagicMock(spec=Exchange)
+    prompt = "New prompt"
+
+    with patch("goose.utils.ask.replace_prompt") as mock_replace_prompt:
+        # Configure the mock to return a new mock object with the same spec
+        modified_exchange = MagicMock(spec=Exchange)
+        mock_replace_prompt.return_value = modified_exchange
+
+        ask_an_ai("Test input", exchange, prompt=prompt, no_history=False)
+
+        # Check if replace_prompt was called correctly
+        mock_replace_prompt.assert_called_once_with(exchange, prompt)
+
+        # Assert that the modified exchange was returned correctly
+        assert mock_replace_prompt.return_value is modified_exchange, "Should return the modified exchange mock"
+
+
+def test_ask_an_ai_exchange_usage():
+    """Test that the exchange adds and processes the message correctly."""
+    exchange = MagicMock(spec=Exchange)
+    input_text = "Test input"
+    message_mock = MagicMock(return_value="Mocked Message")
+
+    with patch("goose.utils.ask.Message.user", new=message_mock):
+        ask_an_ai(input_text, exchange, no_history=False)
+
+        # Assert that Message.user was called with the correct input
+        message_mock.assert_called_once_with(input_text)
+
+        # Assert that exchange.add was called with the mocked message
+        exchange.add.assert_called_once_with("Mocked Message")
+        exchange.reply.assert_called_once()
+
+
+def test_ask_an_ai_return_value():
+    """Test that the function returns the correct reply."""
+    exchange = MagicMock(spec=Exchange)
+    expected_reply = "AI response"
+    exchange.reply.return_value = expected_reply
+    result = ask_an_ai("Test input", exchange, no_history=False)
+    assert result == expected_reply, "Function should return the reply from the exchange."
+
+
+# tests for `clear_exchange`
+def test_clear_exchange_without_tools():
+    """Test clearing messages and checkpoints but not tools."""
+    # Arrange
+    exchange = MagicMock(spec=Exchange)
+
+    # Act
+    new_exchange = clear_exchange(exchange, clear_tools=False)
+
+    # Assert
+    exchange.replace.assert_called_once_with(messages=[], checkpoints=[])
+    assert new_exchange == exchange.replace.return_value, "Should return the modified exchange"
+
+
+def test_clear_exchange_with_tools():
+    """Test clearing messages, checkpoints, and tools."""
+    # Arrange
+    exchange = MagicMock(spec=Exchange)
+
+    # Act
+    new_exchange = clear_exchange(exchange, clear_tools=True)
+
+    # Assert
+    exchange.replace.assert_called_once_with(messages=[], checkpoints=[], tools=())
+    assert new_exchange == exchange.replace.return_value, "Should return the modified exchange with tools cleared"
+
+
+def test_clear_exchange_return_value():
+    """Test that the returned value is a new exchange object."""
+    # Arrange
+    exchange = MagicMock(spec=Exchange)
+    new_exchange_mock = MagicMock(spec=Exchange)
+    exchange.replace.return_value = new_exchange_mock
+
+    # Act
+    new_exchange = clear_exchange(exchange, clear_tools=False)
+
+    # Assert
+    assert new_exchange == new_exchange_mock, "Returned exchange should be the new exchange instance"
+
+
+# tests for `replace_prompt`
+def test_replace_prompt():
+    """Test that the system prompt is correctly replaced."""
+    # Arrange
+    exchange = MagicMock(spec=Exchange)
+    prompt = "New system prompt"
+
+    # Act
+    new_exchange = replace_prompt(exchange, prompt)
+
+    # Assert
+    exchange.replace.assert_called_once_with(system=prompt)
+    assert new_exchange == exchange.replace.return_value, "Should return the modified exchange with the new prompt"
+
+
+def test_replace_prompt_return_value():
+    """Test that the returned value is a new exchange object."""
+    # Arrange
+    exchange = MagicMock(spec=Exchange)
+    expected_new_exchange = MagicMock(spec=Exchange)
+    exchange.replace.return_value = expected_new_exchange
+
+    # Act
+    new_exchange = replace_prompt(exchange, "Another prompt")
+
+    # Assert
+    assert new_exchange == expected_new_exchange, "Returned exchange should be the new exchange instance"
@@ -0,0 +1,192 @@
+from unittest.mock import patch
+
+import pytest
+from goose.utils.file_utils import (
+    create_extensions_list,
+    create_language_weighting,
+)  # Adjust the import path as necessary
+
+
+# tests for `create_extensions_list`
+def test_create_extensions_list_valid_input():
+    """Test with valid input and multiple file extensions."""
+    project_root = "/fake/project/root"
+    max_n = 3
+    files = [
+        "/fake/project/root/file1.py",
+        "/fake/project/root/file2.py",
+        "/fake/project/root/file3.md",
+        "/fake/project/root/file4.md",
+        "/fake/project/root/file5.txt",
+        "/fake/project/root/file6.py",
+        "/fake/project/root/file7.md",
+    ]
+
+    with patch("goose.utils.file_utils.create_file_list", return_value=files):
+        extensions = create_extensions_list(project_root, max_n)
+        assert extensions == [".py", ".md", ".txt"], "Should return the top 3 extensions in the correct order"
+
+
+def test_create_extensions_list_zero_max_n():
+    """Test that a ValueError is raised when max_n is 0."""
+    project_root = "/fake/project/root"
+    max_n = 0
+
+    with pytest.raises(ValueError, match="Number of file extensions must be greater than 0"):
+        create_extensions_list(project_root, max_n)
+
+
+def test_create_extensions_list_no_files():
+    """Test with a project root that contains no files."""
+    project_root = "/fake/project/root"
+    max_n = 3
+
+    with patch("goose.utils.file_utils.create_file_list", return_value=[]):
+        extensions = create_extensions_list(project_root, max_n)
+        assert extensions == [], "Should return an empty list when no files are present"
+
+
+def test_create_extensions_list_fewer_extensions_than_max_n():
+    """Test when there are fewer unique extensions than max_n."""
+    project_root = "/fake/project/root"
+    max_n = 5
+    files = [
+        "/fake/project/root/file1.py",
+        "/fake/project/root/file2.py",
+        "/fake/project/root/file3.md",
+    ]
+
+    with patch("goose.utils.file_utils.create_file_list", return_value=files):
+        extensions = create_extensions_list(project_root, max_n)
+        assert extensions == [".py", ".md"], "Should return all available extensions when fewer than max_n"
+
+
+def test_create_extensions_list_files_without_extensions():
+    """Test that files without extensions are ignored."""
+    project_root = "/fake/project/root"
+    max_n = 3
+    files = [
+        "/fake/project/root/file1",
+        "/fake/project/root/file2.py",
+        "/fake/project/root/file3",
+        "/fake/project/root/file4.md",
+    ]
+
+    with patch("goose.utils.file_utils.create_file_list", return_value=files):
+        extensions = create_extensions_list(project_root, max_n)
+        assert extensions == [".py", ".md"], "Should ignore files without extensions"
+
+
+# tests for `create_language_weighting`
+def test_create_language_weighting_normal_case():
+    """Test the function with multiple files and different sizes."""
+    files = [
+        "/fake/project/file1.py",
+        "/fake/project/file2.py",
+        "/fake/project/file3.md",
+        "/fake/project/file4.txt",
+    ]
+
+    sizes = {
+        "/fake/project/file1.py": 100,
+        "/fake/project/file2.py": 200,
+        "/fake/project/file3.md": 50,
+        "/fake/project/file4.txt": 150,
+    }
+
+    # Mocking os.path.getsize to return different sizes for different files
+    with patch("os.path.getsize") as mock_getsize:
+        mock_getsize.side_effect = lambda file: sizes[file]
+
+        result = create_language_weighting(files)
+
+    total = sum(sizes.values())
+
+    expected_result = {
+        ".py": 300 / total * 100,  # 300 out of 600 total
+        ".txt": 150 / total * 100,  # 150 out of 600 total
+        ".md": 50 / total * 100,  # 50 out of 600 total
+    }
+
+    # Check if the result matches the expected output
+    assert result[".py"] == pytest.approx(expected_result.get(".py"), 0.01)
+    assert result[".txt"] == pytest.approx(expected_result.get(".txt"), 0.01)
+    assert result[".md"] == pytest.approx(expected_result.get(".md"), 0.01)
+
+
+def test_create_language_weighting_no_files():
+    """Test the function when no files are provided."""
+    files = []
+
+    result = create_language_weighting(files)
+    assert result == {}, "Should return an empty dictionary when no files are provided"
+
+
+def test_create_language_weighting_files_without_extensions():
+    """Test the function when files have no extensions."""
+    files = [
+        "/fake/project/file1",
+        "/fake/project/file2",
+    ]
+
+    with patch("os.path.getsize", return_value=100):
+        result = create_language_weighting(files)
+
+    assert result == {}, "Should return an empty dictionary when files have no extensions"
+
+
+def test_create_language_weighting_zero_total_size():
+    """Test the function when all files have a size of 0."""
+    files = [
+        "/fake/project/file1.py",
+        "/fake/project/file2.py",
+    ]
+
+    with patch("os.path.getsize", return_value=0):
+        result = create_language_weighting(files)
+
+    assert result == {".py": 0}
+
+
+def test_create_language_weighting_single_file():
+    """Test the function with a single file."""
+    files = [
+        "/fake/project/file1.py",
+    ]
+
+    with patch("os.path.getsize", return_value=100):
+        result = create_language_weighting(files)
+
+    assert result == {".py": 100.0}, "Should return 100% for the single file's extension"
+
+
+def test_create_language_weighting_mixed_extensions():
+    """Test the function with files of mixed extensions and sizes."""
+    files = [
+        "/fake/project/file1.py",
+        "/fake/project/file2.py",
+        "/fake/project/file3.md",
+        "/fake/project/file4.txt",
+        "/fake/project/file5.md",
+    ]
+
+    with patch("os.path.getsize") as mock_getsize:
+        mock_getsize.side_effect = lambda file: {
+            "/fake/project/file1.py": 100,
+            "/fake/project/file2.py": 100,
+            "/fake/project/file3.md": 200,
+            "/fake/project/file4.txt": 300,
+            "/fake/project/file5.md": 100,
+        }[file]
+
+        result = create_language_weighting(files)
+
+    expected_result = {
+        ".txt": 37.5,  # 300 out of 800 total
+        ".md": 37.5,  # 300 out of 800 total
+        ".py": 25.0,  # 200 out of 800 total
+    }
+
+    assert result[".txt"] == pytest.approx(expected_result.get(".txt"), 0.01)
+    assert result[".md"] == pytest.approx(expected_result.get(".md"), 0.01)
+    assert result[".py"] == pytest.approx(expected_result.get(".py"), 0.01)
@@ -0,0 +1,77 @@
+from pathlib import Path
+
+import pytest
+from exchange import Message
+from goose.utils.session_file import list_sorted_session_files, read_from_file, session_file_exists, write_to_file
+
+
+@pytest.fixture
+def file_path(tmp_path):
+    return tmp_path / "test_file.jsonl"
+
+
+def test_read_write_to_file(file_path):
+    messages = [
+        Message.user("prompt1"),
+        Message.user("prompt2"),
+    ]
+    write_to_file(file_path, messages)
+    assert file_path.exists()
+
+    assert read_from_file(file_path) == messages
+
+
+def test_read_from_file_non_existing_file(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        read_from_file(tmp_path / "no_existing.json")
+
+
+def test_read_from_file_non_jsonl_file(file_path):
+    file_path.write_text("Hello World")
+    with pytest.raises(RuntimeError):
+        read_from_file(file_path)
+
+
+def test_list_sorted_session_files(tmp_path):
+    session_files_directory = tmp_path / "session_files_dir"
+    session_files_directory.mkdir()
+    file_names = ["file1", "file2", "file3"]
+    created_session_files = [create_session_file(session_files_directory, file_name) for file_name in file_names]
+
+    sorted_files = list_sorted_session_files(session_files_directory)
+    assert sorted_files == {
+        "file3": created_session_files[2],
+        "file2": created_session_files[1],
+        "file1": created_session_files[0],
+    }
+
+
+def test_list_sorted_session_without_session_files(tmp_path):
+    session_files_directory = tmp_path / "session_files_dir"
+
+    sorted_files = list_sorted_session_files(session_files_directory)
+    assert sorted_files == {}
+
+
+def test_session_file_exists_return_false_when_directory_does_not_exist(tmp_path):
+    session_files_directory = tmp_path / "session_files_dir"
+    assert not session_file_exists(session_files_directory)
+
+
+def test_session_file_exists_return_false_when_no_session_file_exists(tmp_path):
+    session_files_directory = tmp_path / "session_files_dir"
+    session_files_directory.mkdir()
+    assert not session_file_exists(session_files_directory)
+
+
+def test_session_file_exists_return_true_when_session_file_exists(tmp_path):
+    session_files_directory = tmp_path / "session_files_dir"
+    session_files_directory.mkdir()
+    create_session_file(session_files_directory, "session1")
+    assert session_file_exists(session_files_directory)
+
+
+def create_session_file(file_path, file_name) -> Path:
+    file = file_path / f"{file_name}.jsonl"
+    file.touch()
+    return file
@@ -0,0 +1,63 @@
+import string
+
+import pytest
+from goose.utils import droid, ensure, ensure_list, load_plugins
+
+
+class MockClass:
+    def __init__(self, name):
+        self.name = name
+
+    def __eq__(self, other):
+        return self.name == other.name
+
+
+def test_load_plugins():
+    plugins = load_plugins("exchange.provider")
+    assert isinstance(plugins, dict)
+    assert len(plugins) > 0
+
+
+def test_ensure_with_class():
+    mock_class = MockClass("foo")
+    assert ensure(MockClass)(mock_class) == mock_class
+
+
+def test_ensure_with_dictionary():
+    mock_class = ensure(MockClass)({"name": "foo"})
+    assert mock_class == MockClass("foo")
+
+
+def test_ensure_with_invalid_dictionary():
+    with pytest.raises(TypeError):
+        ensure(MockClass)({"age": "foo"})
+
+
+def test_ensure_with_list():
+    mock_class = ensure(MockClass)(["foo"])
+    assert mock_class == MockClass("foo")
+
+
+def test_ensure_with_invalid_list():
+    with pytest.raises(TypeError):
+        ensure(MockClass)(["foo", "bar"])
+
+
+def test_ensure_with_value():
+    mock_class = ensure(MockClass)("foo")
+    assert mock_class == MockClass("foo")
+
+
+def test_ensure_list():
+    obj_list = ensure_list(MockClass)(["foo", "bar"])
+    assert obj_list == [MockClass("foo"), MockClass("bar")]
+
+
+def test_droid():
+    result = droid()
+    assert isinstance(result, str)
+    assert len(result) == 4
+    for character in [result[i] for i in [0, 2]]:
+        assert character in string.ascii_lowercase, "should be in lower case"
+    for character in [result[i] for i in [1, 3]]:
+        assert character in string.digits, "should be a digit"