Fix tests

tests/benchmark/test_report_v2.py

@@ -5,6 +5,7 @@
 import pytest
 
 from moatless.benchmark.report_v2 import EditStats, PlanStats, SearchStats, StateStats, to_result, BenchmarkResult, to_dataframe
+from moatless.benchmark.utils import get_moatless_instance
 from moatless.trajectory import Trajectory
 
 
@@ -13,6 +14,11 @@ def django_trajectory():
     file_path = Path("tests/trajectories/django__django_16379.json")
     return Trajectory.load(str(file_path))
 
+@pytest.fixture
+def scikit_trajectory():
+    file_path = Path("tests/trajectories/scikit-learn__scikit-learn-13779/trajectory.json")
+    return Trajectory.load(str(file_path))
+
 
 @pytest.fixture
 def dataset():
@@ -22,10 +28,12 @@ def dataset():
 
 @pytest.fixture
 def django_instance(dataset):
-    for instance in dataset:
-        if instance["instance_id"] == "django__django-16379":
-            return instance
-    return None
+    return get_moatless_instance("django__django-16379", split="lite")
+
+
+@pytest.fixture
+def scikit_instance(dataset):
+    return get_moatless_instance("scikit-learn__scikit-learn-13779", split="verified")
 
 
 def test_to_result(django_trajectory, django_instance):
@@ -69,6 +77,14 @@ def test_to_result(django_trajectory, django_instance):
     assert result.expected_files == len(result.expected_spans_details)
 
 
+def test_scikit_not_edited(scikit_trajectory, scikit_instance):
+    result = to_result(scikit_instance, scikit_trajectory)
+
+    print(json.dumps(result.model_dump(), indent=2))
+
+    assert result.edit.status == "expected_files"
+
+
 def test_to_result_error_case(django_trajectory, django_instance):
     # Simulate an error in the trajectory
     django_trajectory._info["error"] = "Simulated error"

tests/edit/test_clarify.py

@@ -1,7 +1,10 @@
 import pytest
 from unittest.mock import Mock, patch
+
+from moatless.benchmark.swebench import create_workspace
+from moatless.benchmark.utils import get_moatless_instance
 from moatless.edit.clarify import ClarifyCodeChange, LineNumberClarification
-from moatless.schema import StateOutcome, FileWithSpans
+from moatless.state import StateOutcome, FileWithSpans
 from moatless.workspace import Workspace
 from moatless.file_context import FileContext
 from moatless.repository import CodeFile
@@ -133,3 +136,15 @@ def test_verify_line_numbers_invalid(self, mock_span, mock_file, clarify_code_ch
 
         assert result is not None
         assert "covers the whole code span" in result
+
+def test_clarify_class():
+    instance = get_moatless_instance("django__django-11095", split="verified")
+    workspace = create_workspace(instance)
+    workspace.file_context.add_spans_to_context("django/contrib/admin/checks.py", ["ModelAdminChecks", "ModelAdminChecks._check_inlines_item"])
+
+    clarify_code = ClarifyCodeChange(id=0, _workspace=workspace, initial_message="Test initial message", instructions="update", span_id="ModelAdminChecks", file_path="django/contrib/admin/checks.py")
+    prompt_message = clarify_code.messages()[-1].content
+
+    assert "class ModelAdminChecks(BaseModelAdminChecks):" in prompt_message
+    assert "def _check_inlines_item(self, obj, inline, label):" in prompt_message
+    assert "return inline(obj.model, obj.admin_site).check()" in prompt_message

tests/edit/test_edit.py

@@ -5,7 +5,7 @@
 from moatless.benchmark.swebench import load_instance, create_workspace
 from moatless.edit.edit import EditCode
 from moatless.repository.file import UpdateResult
-from moatless.schema import StateOutcome, Content
+from moatless.state import StateOutcome, Content
 from moatless.settings import Settings
 from moatless.trajectory import Trajectory
 from moatless.workspace import Workspace

tests/edit/test_plan.py

@@ -1,7 +1,11 @@
 import pytest
 from unittest.mock import Mock, patch
+
+from moatless.benchmark.swebench import create_workspace
+from moatless.benchmark.utils import get_moatless_instance
+from moatless.edit.expand import ExpandContext
 from moatless.edit.plan import PlanToCode, ApplyChange
-from moatless.schema import StateOutcome, ActionTransaction
+from moatless.state import StateOutcome, ActionTransaction
 from moatless.workspace import Workspace
 from moatless.file_context import FileContext
 
@@ -112,4 +116,46 @@ def test_request_for_change_file_not_found(self, mock_get_spans, mock_get_file,
         assert response.trigger == "retry"
         assert "File nonexistent.py is not found in the file context" in response.retry_message
 
-    # Add more tests for other scenarios in _request_for_change method
+def test_select_class():
+    instance = get_moatless_instance("django__django-11095", split="verified")
+    workspace = create_workspace(instance)
+
+    plan_to_code = PlanToCode(id=0, _workspace=workspace, initial_message="Test initial message")
+
+    workspace.file_context.add_spans_to_context("django/contrib/admin/checks.py", ["ModelAdminChecks", "ModelAdminChecks._check_inlines_item"])
+    mocked_action = ApplyChange(
+        scratch_pad="Applying change",
+        action="modify",
+        file_path="django/contrib/admin/checks.py",
+        span_id="ModelAdminChecks",
+        instructions="Update function"
+    )
+
+    outcome = plan_to_code.execute(mocked_action)
+
+    # Expect clarification
+    assert outcome.trigger == "edit_code"
+    assert outcome.output == {'instructions': 'Update function', 'file_path': 'django/contrib/admin/checks.py', 'span_id': 'ModelAdminChecks'}
+
+
+def test_impl_span():
+    instance = get_moatless_instance("django__django-12419", split="verified")
+    workspace = create_workspace(instance)
+
+    plan_to_code = PlanToCode(id=0, _workspace=workspace, initial_message="Test initial message")
+
+    workspace.file_context.add_spans_to_context("django/contrib/admin/checks.py", ["ModelAdminChecks", "ModelAdminChecks._check_inlines_item"])
+    mocked_action = ApplyChange(
+        scratch_pad="Applying change",
+        action="modify",
+        file_path="django/contrib/admin/checks.py",
+        span_id="ModelAdminChecks",
+        instructions="Update function"
+    )
+
+    outcome = plan_to_code.execute(mocked_action)
+
+    # Expect clarification
+    assert outcome.trigger == "edit_code"
+    assert outcome.output == {'instructions': 'Update function', 'file_path': 'django/contrib/admin/checks.py', 'span_id': 'ModelAdminChecks'}
+

tests/find/test_decide.py

@@ -1,7 +1,7 @@
 import pytest
 from moatless.find.decide import DecideRelevance, Decision
 from moatless.find.identify import Identify, IdentifyCode
-from moatless.schema import StateOutcome, ActionTransaction
+from moatless.state import StateOutcome, ActionTransaction
 from moatless.workspace import Workspace
 from moatless.file_context import FileContext
 from unittest.mock import Mock, MagicMock, patch

tests/find/test_identify.py

@@ -3,7 +3,8 @@
 from moatless.find.identify import IdentifyCode, Identify, is_test_pattern
 from moatless.file_context import RankedFileSpan
 from moatless.repository.file import CodeFile
-from moatless.schema import FileWithSpans, StateOutcome
+from moatless.schema import FileWithSpans
+from moatless.state import StateOutcome
 from moatless.workspace import Workspace
 from unittest.mock import Mock, MagicMock
 

tests/find/test_search.py

@@ -1,6 +1,9 @@
 import pytest
+
+from moatless.benchmark.swebench import create_workspace
+from moatless.benchmark.utils import get_moatless_instance, get_moatless_instances
 from moatless.find.search import SearchCode, Search, SearchRequest
-from moatless.schema import StateOutcome
+from moatless.state import StateOutcome
 from moatless.workspace import Workspace
 from unittest.mock import Mock, MagicMock
 from pydantic import ValidationError
@@ -117,4 +120,82 @@ def test_handle_direct_search_attributes(self):
 
         assert len(search.search_requests) == 1
         assert search.search_requests[0].file_pattern == "*.js"
-        assert search.search_requests[0].query == "javascript query"
+        assert search.search_requests[0].query == "javascript query"
+
+
+def test_find_impl_span():
+
+    instances = get_moatless_instances(split="verified")
+
+    # Filter and sort instances
+    filtered_instances = {
+        k: v for k, v in instances.items()
+        if "django__django-" in k and "12273" <= k.split("-")[-1] <= "12419"
+    }
+    sorted_instances = dict(sorted(filtered_instances.items()))
+
+    for instance_id, instance in sorted_instances.items():
+        print(f"Instance: {instance_id}")
+        workspace = create_workspace(instance)
+
+        search_code = SearchCode(id=0, _workspace=workspace, initial_message="Test initial message")
+
+        mocked_action = Search(
+            scratch_pad="Applying change",
+            search_requests=[
+                SearchRequest(file_pattern="**/global_settings.py", query="SECURE_REFERRER_POLICY setting")
+            ]
+        )
+
+        outcome = search_code.execute(mocked_action)
+        print(outcome)
+
+        workspace.file_context.add_ranked_spans(outcome.output["ranked_spans"])
+        assert "SECURE_REFERRER_POLICY" in workspace.file_context.create_prompt()
+
+def test_find():
+    instance_id = "django__django-12419" #
+    instance = get_moatless_instance(instance_id, split="verified")
+    print(f"Instance: {instance_id}")
+    workspace = create_workspace(instance)
+
+    search_code = SearchCode(id=0, _workspace=workspace, initial_message="Test initial message")
+
+    mocked_action = Search(
+        scratch_pad="Applying change",
+        search_requests=[
+            SearchRequest(file_pattern="**/global_settings.py", query="SECURE_REFERRER_POLICY setting")
+        ]
+    )
+
+    outcome = search_code.execute(mocked_action)
+
+    for ranked_span in outcome.output["ranked_spans"]:
+        print(ranked_span)
+
+    workspace.file_context.add_ranked_spans(outcome.output["ranked_spans"])
+    print(workspace.file_context.create_prompt(show_span_ids=True))
+    assert "SECURE_REFERRER_POLICY = None" in workspace.file_context.create_prompt()
+
+
+def test_find_2():
+    instance_id = "django__django-15104"
+    instance = get_moatless_instance(instance_id, split="verified")
+    workspace = create_workspace(instance)
+
+    search_code = SearchCode(id=0, _workspace=workspace, initial_message="Test initial message")
+
+    print(instance["expected_spans"])
+
+    mocked_action = Search(
+        scratch_pad="Applying change",
+        max_search_results=250,
+        search_requests=[
+            SearchRequest(file_pattern="**/migrations/*.py", query="MigrationAutodetector class with generate_renamed_models method")
+        ]
+    )
+
+    outcome = search_code.execute(mocked_action)
+
+    for span in outcome.output["ranked_spans"]:
+        print(span)

tests/index/test_code_index.py

@@ -0,0 +1,32 @@
+from moatless.benchmark.swebench import setup_swebench_repo
+from moatless.benchmark.utils import get_moatless_instance
+from moatless.index import IndexSettings, CodeIndex
+from moatless.index.settings import CommentStrategy
+from moatless.repository import GitRepository, FileRepository
+
+
+def test_ingestion():
+    index_settings = IndexSettings(
+        embed_model="voyage-code-2",
+        dimensions=1536,
+        language="python",
+        min_chunk_size=200,
+        chunk_size=750,
+        hard_token_limit=3000,
+        max_chunks=200,
+        comment_strategy=CommentStrategy.ASSOCIATE,
+    )
+
+    instance_id = "django__django-12419"
+    instance = get_moatless_instance(instance_id, split="verified")
+    repo_dir = setup_swebench_repo(instance)
+    print(repo_dir)
+    repo = FileRepository(repo_dir)
+    code_index = CodeIndex(settings=index_settings, file_repo=repo)
+
+    vectors, indexed_tokens = code_index.run_ingestion(num_workers=1, input_files=["django/conf/global_settings.py"])
+
+    results = code_index._vector_search("SECURE_REFERRER_POLICY setting")
+
+    for result in results:
+        print(result)

tests/index/test_epic_split.py

@@ -3,7 +3,8 @@
 
 from llama_index.core import SimpleDirectoryReader
 
-from moatless.benchmark.swebench import get_moatless_instance, setup_swebench_repo
+from moatless.benchmark.swebench import setup_swebench_repo
+from moatless.benchmark.utils import get_moatless_instance
 from moatless.index import IndexSettings
 from moatless.index.epic_split import EpicSplitter
 
@@ -48,5 +49,49 @@ def test_epic_split():
     stats.print_stats(40)
 
 
+    for node in prepared_nodes:
+        print(f"{node.id_} {node.metadata['tokens']} {node.metadata['start_line']}-{node.metadata['end_line']} {node.metadata['span_ids']}")
+
+
+def test_impl_spans():
+    instance_id = "django__django-12419"
+    instance = get_moatless_instance(instance_id, split="verified")
+    repo_path = setup_swebench_repo(instance)
+
+    file = "django/conf/global_settings.py"
+    input_files = [f"{repo_path}/{file}"]
+
+    settings = IndexSettings()
+    splitter = EpicSplitter(
+        language=settings.language,
+        min_chunk_size=settings.min_chunk_size,
+        chunk_size=settings.chunk_size,
+        hard_token_limit=settings.hard_token_limit,
+        max_chunks=settings.max_chunks,
+        comment_strategy=settings.comment_strategy,
+        min_lines_to_parse_block=50,
+        repo_path=repo_path,
+    )
+
+    reader = SimpleDirectoryReader(
+        input_dir=repo_path,
+        input_files=input_files,
+        filename_as_id=True,
+        recursive=True,
+    )
+    docs = reader.load_data()
+    print(f"Read {len(docs)} documents")
+
+    # Profile the get_nodes_from_documents method
+    profiler = cProfile.Profile()
+    profiler.enable()
+    prepared_nodes = splitter.get_nodes_from_documents(docs, show_progress=True)
+    profiler.disable()
+
+    # Print the profiling results
+    stats = pstats.Stats(profiler).sort_stats(pstats.SortKey.CUMULATIVE)
+    stats.print_stats(40)
+
+
     for node in prepared_nodes:
         print(f"{node.id_} {node.metadata['tokens']} {node.metadata['start_line']}-{node.metadata['end_line']} {node.metadata['span_ids']}")

tests/index/test_simple_faiss.py

@@ -0,0 +1,19 @@
+import os
+
+from llama_index.core.vector_stores import VectorStoreQuery
+
+from moatless.index import SimpleFaissVectorStore
+
+
+def test_search_index():
+    index_store_dir = os.getenv("INDEX_STORE_DIR")
+    vector_store = SimpleFaissVectorStore.from_persist_dir(f"{index_store_dir}/django__django-12419")
+
+    query_bundle = VectorStoreQuery(
+        query_str="SECURE_REFERRER_POLICY setting",
+        similarity_top_k=100,
+    )
+
+    result = vector_store.query(query_bundle)
+    for res in result.ids:
+        print(res)