diff --git a/mock_tests/test_batch.py b/mock_tests/test_batch.py index c00bd08db..ad7357f04 100644 --- a/mock_tests/test_batch.py +++ b/mock_tests/test_batch.py @@ -1,7 +1,8 @@ -from typing import Generator +from typing import AsyncGenerator, Generator, List import grpc import pytest +import pytest_asyncio import weaviate from weaviate.proto.v1 import batch_pb2, weaviate_pb2_grpc from .conftest import MOCK_IP, MOCK_PORT, MOCK_PORT_GRPC, mock_class, HTTPServer @@ -60,3 +61,102 @@ def test_ssb_canceled_stream( for i in range(HOW_MANY): batch.add_object({"name": f"Object {i}"}) assert len(service.uuids) == HOW_MANY + + +class MockFailedObjectWeaviateService(weaviate_pb2_grpc.WeaviateServicer): + """Rejects every other object, starting with the first. + + A batch of 1 object gives 1 error and 0 uuids; a batch of 4 gives 2 errors and 2 uuids. + """ + + def __init__(self) -> None: + self.seen = 0 + + def BatchStream( + self, + request_iterator: Generator[batch_pb2.BatchStreamRequest, None, None], + context: grpc.ServicerContext, + ) -> Generator[batch_pb2.BatchStreamReply, None, None]: + yield batch_pb2.BatchStreamReply(started=batch_pb2.BatchStreamReply.Started()) + for request in request_iterator: + if request.HasField("data"): + uuids: List[str] = [] + errors: List[batch_pb2.BatchStreamReply.Results.Error] = [] + successes: List[batch_pb2.BatchStreamReply.Results.Success] = [] + for obj in request.data.objects.values: + uuids.append(obj.uuid) + if self.seen % 2 == 0: + errors.append( + batch_pb2.BatchStreamReply.Results.Error( + uuid=obj.uuid, error="mock failure" + ) + ) + else: + successes.append(batch_pb2.BatchStreamReply.Results.Success(uuid=obj.uuid)) + self.seen += 1 + yield batch_pb2.BatchStreamReply(acks=batch_pb2.BatchStreamReply.Acks(uuids=uuids)) + yield batch_pb2.BatchStreamReply( + results=batch_pb2.BatchStreamReply.Results(errors=errors, successes=successes) + ) + if request.HasField("stop"): + return + + +@pytest.fixture(scope="function") +def failed_object_stream( + canceled_stream_client: weaviate.WeaviateClient, start_grpc_server: grpc.Server +): + service = MockFailedObjectWeaviateService() + weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server) + return canceled_stream_client.collections.use(mock_class["class"]) + + +@pytest_asyncio.fixture +async def failed_object_stream_async( + weaviate_mock: HTTPServer, start_grpc_server: grpc.Server +) -> AsyncGenerator[weaviate.collections.CollectionAsync, None]: + weaviate_mock.expect_request(f"/v1/schema/{mock_class['class']}").respond_with_json(mock_class) + weaviate_pb2_grpc.add_WeaviateServicer_to_server( + MockFailedObjectWeaviateService(), start_grpc_server + ) + client = weaviate.use_async_with_local(port=MOCK_PORT, host=MOCK_IP, grpc_port=MOCK_PORT_GRPC) + await client.connect() + yield client.collections.use(mock_class["class"]) + await client.close() + + +def test_ingest_has_errors_on_failed_object( + failed_object_stream: weaviate.collections.Collection, +): + result = failed_object_stream.data.ingest([{"name": "Object 1"}]) + assert result.has_errors is True + assert len(result.errors) == 1 + + +def test_ssb_ingest_reports_has_errors( + failed_object_stream: weaviate.collections.Collection, +) -> None: + result = failed_object_stream.data.ingest({"name": f"Object {i}"} for i in range(4)) + assert len(result.errors) == 2 + assert len(result.uuids) == 2 + assert result.has_errors + + +@pytest.mark.asyncio +async def test_ssb_ingest_reports_has_errors_async( + failed_object_stream_async: weaviate.collections.CollectionAsync, +) -> None: + result = await failed_object_stream_async.data.ingest({"name": f"Object {i}"} for i in range(4)) + assert len(result.errors) == 2 + assert len(result.uuids) == 2 + assert result.has_errors + + +def test_ssb_stream_reports_has_errors( + failed_object_stream: weaviate.collections.Collection, +) -> None: + with failed_object_stream.batch.stream() as batch: + for i in range(4): + batch.add_object({"name": f"Object {i}"}) + assert len(failed_object_stream.batch.failed_objects) == 2 + assert failed_object_stream.batch.results.objs.has_errors diff --git a/test/collection/test_batch.py b/test/collection/test_batch.py index 0a2cda954..3837317ab 100644 --- a/test/collection/test_batch.py +++ b/test/collection/test_batch.py @@ -3,10 +3,38 @@ import pytest from weaviate.collections.batch.grpc_batch import _validate_props -from weaviate.collections.classes.batch import MAX_STORED_RESULTS, BatchObjectReturn +from weaviate.collections.classes.batch import ( + MAX_STORED_RESULTS, + BatchObject, + BatchObjectReturn, + BatchReference, + BatchReferenceReturn, + ErrorObject, + ErrorReference, +) from weaviate.exceptions import WeaviateInsertInvalidPropertyError +def _error_object(index: int) -> ErrorObject: + return ErrorObject( + message="something went wrong", + object_=BatchObject(collection="Test", properties={"name": "test"}, index=index), + ) + + +def _error_reference(index: int) -> ErrorReference: + return ErrorReference( + message="something went wrong", + reference=BatchReference( + from_object_collection="Test", + from_object_uuid=uuid.uuid4(), + from_property_name="other", + to_object_uuid=uuid.uuid4(), + index=index, + ), + ) + + def test_batch_object_return_add() -> None: lhs_uuids = [uuid.uuid4() for _ in range(MAX_STORED_RESULTS)] lhs = BatchObjectReturn( @@ -36,6 +64,42 @@ def test_batch_object_return_add() -> None: } +def test_batch_object_return_has_errors_when_constructed_with_errors() -> None: + err = _error_object(0) + result = BatchObjectReturn(_all_responses=[err], errors={0: err}) + assert result.has_errors + + +def test_batch_object_return_add_sets_has_errors() -> None: + err = _error_object(1) + result = BatchObjectReturn() + result += BatchObjectReturn(_all_responses=[uuid.uuid4()], uuids={0: uuid.uuid4()}) + result += BatchObjectReturn(_all_responses=[err], errors={1: err}) + assert result.has_errors + assert len(result.errors) == 1 + + +def test_batch_object_return_has_no_errors_when_all_succeed() -> None: + uid = uuid.uuid4() + result = BatchObjectReturn() + result += BatchObjectReturn(_all_responses=[uid], uuids={0: uid}) + assert not result.has_errors + + +def test_batch_reference_return_has_errors_when_constructed_with_errors() -> None: + err = _error_reference(0) + result = BatchReferenceReturn(errors={0: err}) + assert result.has_errors + + +def test_batch_reference_return_add_sets_has_errors() -> None: + err = _error_reference(0) + result = BatchReferenceReturn() + result += BatchReferenceReturn(errors={0: err}) + assert result.has_errors + assert len(result.errors) == 1 + + def test_validate_props_raises_for_top_level_id() -> None: with pytest.raises(WeaviateInsertInvalidPropertyError): _validate_props({"id": "abc123"}) diff --git a/weaviate/collections/classes/batch.py b/weaviate/collections/classes/batch.py index eb8d0181a..b90bebfad 100644 --- a/weaviate/collections/classes/batch.py +++ b/weaviate/collections/classes/batch.py @@ -216,6 +216,9 @@ class BatchObjectReturn: uuids: Dict[int, uuid_package.UUID] = field(default_factory=dict) has_errors: bool = False + def __post_init__(self) -> None: + self.has_errors = self.has_errors or len(self.errors) > 0 + @property def all_responses(self) -> List[Union[uuid_package.UUID, ErrorObject]]: """@deprecated: A list of all the responses from the batch operation. Each response is either a `uuid_package.UUID` object or an `Error` object. @@ -284,6 +287,9 @@ class BatchReferenceReturn: errors: Dict[int, ErrorReference] = field(default_factory=dict) has_errors: bool = False + def __post_init__(self) -> None: + self.has_errors = self.has_errors or len(self.errors) > 0 + def __add__(self, other: "BatchReferenceReturn") -> "BatchReferenceReturn": self.elapsed_seconds += other.elapsed_seconds prev_max = max(self.errors.keys()) if len(self.errors) > 0 else -1