-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'main' of https://github.com/onyx-dot-app/onyx into feat…
…ure/more_celery_fanout
- Loading branch information
Showing
10 changed files
with
267 additions
and
23 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
106 changes: 106 additions & 0 deletions
106
backend/tests/integration/common_utils/managers/index_attempt.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,106 @@ | ||
from datetime import datetime | ||
from datetime import timedelta | ||
from urllib.parse import urlencode | ||
|
||
import requests | ||
|
||
from onyx.db.engine import get_session_context_manager | ||
from onyx.db.enums import IndexModelStatus | ||
from onyx.db.models import IndexAttempt | ||
from onyx.db.models import IndexingStatus | ||
from onyx.db.search_settings import get_current_search_settings | ||
from onyx.server.documents.models import IndexAttemptSnapshot | ||
from onyx.server.documents.models import PaginatedReturn | ||
from tests.integration.common_utils.constants import API_SERVER_URL | ||
from tests.integration.common_utils.constants import GENERAL_HEADERS | ||
from tests.integration.common_utils.test_models import DATestIndexAttempt | ||
from tests.integration.common_utils.test_models import DATestUser | ||
|
||
|
||
class IndexAttemptManager: | ||
@staticmethod | ||
def create_test_index_attempts( | ||
num_attempts: int, | ||
cc_pair_id: int, | ||
from_beginning: bool = False, | ||
status: IndexingStatus = IndexingStatus.SUCCESS, | ||
new_docs_indexed: int = 10, | ||
total_docs_indexed: int = 10, | ||
docs_removed_from_index: int = 0, | ||
error_msg: str | None = None, | ||
base_time: datetime | None = None, | ||
) -> list[DATestIndexAttempt]: | ||
if base_time is None: | ||
base_time = datetime.now() | ||
|
||
attempts = [] | ||
with get_session_context_manager() as db_session: | ||
# Get the current search settings | ||
search_settings = get_current_search_settings(db_session) | ||
if ( | ||
not search_settings | ||
or search_settings.status != IndexModelStatus.PRESENT | ||
): | ||
raise ValueError("No current search settings found with PRESENT status") | ||
|
||
for i in range(num_attempts): | ||
time_created = base_time - timedelta(hours=i) | ||
|
||
index_attempt = IndexAttempt( | ||
connector_credential_pair_id=cc_pair_id, | ||
from_beginning=from_beginning, | ||
status=status, | ||
new_docs_indexed=new_docs_indexed, | ||
total_docs_indexed=total_docs_indexed, | ||
docs_removed_from_index=docs_removed_from_index, | ||
error_msg=error_msg, | ||
time_created=time_created, | ||
time_started=time_created, | ||
time_updated=time_created, | ||
search_settings_id=search_settings.id, | ||
) | ||
|
||
db_session.add(index_attempt) | ||
db_session.flush() # To get the ID | ||
|
||
attempts.append( | ||
DATestIndexAttempt( | ||
id=index_attempt.id, | ||
status=index_attempt.status, | ||
new_docs_indexed=index_attempt.new_docs_indexed, | ||
total_docs_indexed=index_attempt.total_docs_indexed, | ||
docs_removed_from_index=index_attempt.docs_removed_from_index, | ||
error_msg=index_attempt.error_msg, | ||
time_started=index_attempt.time_started, | ||
time_updated=index_attempt.time_updated, | ||
) | ||
) | ||
|
||
db_session.commit() | ||
|
||
return attempts | ||
|
||
@staticmethod | ||
def get_index_attempt_page( | ||
cc_pair_id: int, | ||
page: int = 0, | ||
page_size: int = 10, | ||
user_performing_action: DATestUser | None = None, | ||
) -> PaginatedReturn[IndexAttemptSnapshot]: | ||
query_params: dict[str, str | int] = { | ||
"page_num": page, | ||
"page_size": page_size, | ||
} | ||
|
||
response = requests.get( | ||
url=f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}/index-attempts?{urlencode(query_params, doseq=True)}", | ||
headers=user_performing_action.headers | ||
if user_performing_action | ||
else GENERAL_HEADERS, | ||
) | ||
response.raise_for_status() | ||
data = response.json() | ||
return PaginatedReturn( | ||
items=[IndexAttemptSnapshot(**item) for item in data["items"]], | ||
total_items=data["total_items"], | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
90 changes: 90 additions & 0 deletions
90
backend/tests/integration/tests/index_attempt/test_index_attempt_pagination.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,90 @@ | ||
from datetime import datetime | ||
|
||
from onyx.db.models import IndexingStatus | ||
from tests.integration.common_utils.managers.cc_pair import CCPairManager | ||
from tests.integration.common_utils.managers.index_attempt import IndexAttemptManager | ||
from tests.integration.common_utils.managers.user import UserManager | ||
from tests.integration.common_utils.test_models import DATestIndexAttempt | ||
from tests.integration.common_utils.test_models import DATestUser | ||
|
||
|
||
def _verify_index_attempt_pagination( | ||
cc_pair_id: int, | ||
index_attempts: list[DATestIndexAttempt], | ||
page_size: int = 5, | ||
user_performing_action: DATestUser | None = None, | ||
) -> None: | ||
retrieved_attempts: list[int] = [] | ||
last_time_started = None # Track the last time_started seen | ||
|
||
for i in range(0, len(index_attempts), page_size): | ||
paginated_result = IndexAttemptManager.get_index_attempt_page( | ||
cc_pair_id=cc_pair_id, | ||
page=(i // page_size), | ||
page_size=page_size, | ||
user_performing_action=user_performing_action, | ||
) | ||
|
||
# Verify that the total items is equal to the length of the index attempts list | ||
assert paginated_result.total_items == len(index_attempts) | ||
# Verify that the number of items in the page is equal to the page size | ||
assert len(paginated_result.items) == min(page_size, len(index_attempts) - i) | ||
|
||
# Verify time ordering within the page (descending order) | ||
for attempt in paginated_result.items: | ||
if last_time_started is not None: | ||
assert ( | ||
attempt.time_started <= last_time_started | ||
), "Index attempts not in descending time order" | ||
last_time_started = attempt.time_started | ||
|
||
# Add the retrieved index attempts to the list of retrieved attempts | ||
retrieved_attempts.extend([attempt.id for attempt in paginated_result.items]) | ||
|
||
# Create a set of all the expected index attempt IDs | ||
all_expected_attempts = set(attempt.id for attempt in index_attempts) | ||
# Create a set of all the retrieved index attempt IDs | ||
all_retrieved_attempts = set(retrieved_attempts) | ||
|
||
# Verify that the set of retrieved attempts is equal to the set of expected attempts | ||
assert all_expected_attempts == all_retrieved_attempts | ||
|
||
|
||
def test_index_attempt_pagination(reset: None) -> None: | ||
# Create an admin user to perform actions | ||
user_performing_action: DATestUser = UserManager.create( | ||
name="admin_performing_action", | ||
is_first_user=True, | ||
) | ||
|
||
# Create a CC pair to attach index attempts to | ||
cc_pair = CCPairManager.create_from_scratch( | ||
user_performing_action=user_performing_action, | ||
) | ||
|
||
# Create 300 successful index attempts | ||
base_time = datetime.now() | ||
all_attempts = IndexAttemptManager.create_test_index_attempts( | ||
num_attempts=300, | ||
cc_pair_id=cc_pair.id, | ||
status=IndexingStatus.SUCCESS, | ||
base_time=base_time, | ||
) | ||
|
||
# Verify basic pagination with different page sizes | ||
print("Verifying basic pagination with page size 5") | ||
_verify_index_attempt_pagination( | ||
cc_pair_id=cc_pair.id, | ||
index_attempts=all_attempts, | ||
page_size=5, | ||
user_performing_action=user_performing_action, | ||
) | ||
|
||
# Test with a larger page size | ||
print("Verifying pagination with page size 100") | ||
_verify_index_attempt_pagination( | ||
cc_pair_id=cc_pair.id, | ||
index_attempts=all_attempts, | ||
page_size=100, | ||
user_performing_action=user_performing_action, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.