chore: fix flaky unit test and add proper shutdown for file batches (#3725)

# What does this PR do?
Have been running into flaky unit test failures:
5217035494
Fixing below
1. Shutting down properly by cancelling any stale file batches tasks
running in background.
2. Also, use unique_kvstore_config, so the test dont use same db path
and maintain test isolation
## Test Plan
Ran unit test locally and CI
This commit is contained in:
slekkala1 2025-10-07 14:23:14 -07:00 committed by GitHub
parent 1970b4aa4b
commit c2d97a9db9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 39 additions and 17 deletions

View file

@ -293,6 +293,19 @@ class OpenAIVectorStoreMixin(ABC):
await self._resume_incomplete_batches()
self._last_file_batch_cleanup_time = 0
async def shutdown(self) -> None:
"""Clean up mixin resources including background tasks."""
# Cancel any running file batch tasks gracefully
if hasattr(self, "_file_batch_tasks"):
tasks_to_cancel = list(self._file_batch_tasks.items())
for _, task in tasks_to_cancel:
if not task.done():
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
@abstractmethod
async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
"""Delete chunks from a vector store."""