diff --git a/chromadb/test/property/test_persist.py b/chromadb/test/property/test_persist.py index 92f65b27714..907ce94459a 100644 --- a/chromadb/test/property/test_persist.py +++ b/chromadb/test/property/test_persist.py @@ -11,7 +11,8 @@ import chromadb from chromadb.api import ClientAPI, ServerAPI from chromadb.config import Settings, System -from chromadb.segment import SegmentManager, VectorReader +from chromadb.segment import SegmentManager +from chromadb.segment.impl.vector.local_hnsw import LocalHnswSegment import chromadb.test.property.strategies as strategies import chromadb.test.property.invariants as invariants from chromadb.test.property.test_embeddings import ( @@ -142,7 +143,7 @@ def test_sync_threshold(settings: Settings) -> None: ) manager = system.instance(SegmentManager) - segment = manager.get_segment(collection.id, VectorReader) + segment = manager.get_segment(collection.id, LocalHnswSegment) def get_index_last_modified_at() -> float: # Time resolution on Windows can be up to 10ms diff --git a/chromadb/test/segment/test_vector.py b/chromadb/test/segment/test_vector.py index 0d62c827461..07c0b61d709 100644 --- a/chromadb/test/segment/test_vector.py +++ b/chromadb/test/segment/test_vector.py @@ -1,5 +1,5 @@ import pytest -from typing import Generator, List, Callable, Iterator, Type, cast +from typing import Generator, List, Callable, Iterator, Type, cast, Union from chromadb.config import System, Settings from chromadb.test.conftest import ProducerFn from chromadb.types import ( @@ -14,7 +14,6 @@ Vector, ) from chromadb.ingest import Producer -from chromadb.segment import VectorReader import uuid import time @@ -34,6 +33,8 @@ import shutil import numpy as np +VectorReader = Union[LocalHnswSegment, PersistentLocalHnswSegment] + def sqlite() -> Generator[System, None, None]: """Fixture generator for sqlite DB"""