Skip to content

Commit

Permalink
fix: use correct path to reset RAG manager state between tests
Browse files Browse the repository at this point in the history
- Fix type error by using correct internal path to _rag_manager
- Split imports for better readability
- Organize imports more clearly
  • Loading branch information
ErikBjare committed Nov 22, 2024
1 parent 0e9cb11 commit edbf520
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions tests/test_tools_rag.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,24 @@

from unittest.mock import patch

import gptme.tools._rag_context
import gptme.tools.rag
import pytest
from gptme.message import Message
from gptme.tools._rag_context import enhance_messages
from gptme.tools.rag import _HAS_RAG, init as init_rag, rag_index, rag_search
from gptme.tools.rag import _HAS_RAG
from gptme.tools.rag import init as init_rag
from gptme.tools.rag import rag_index, rag_search


@pytest.fixture(autouse=True)
def reset_rag():
"""Reset the RAG manager and init state before and after each test."""
import gptme.tools.rag

gptme.tools.rag.rag_manager = None
gptme.tools._rag_context._rag_manager = None
gptme.tools.rag._init_run = False
yield
gptme.tools.rag.rag_manager = None
gptme.tools._rag_context._rag_manager = None
gptme.tools.rag._init_run = False


Expand Down

0 comments on commit edbf520

Please sign in to comment.