def test_write(self): # Setup Expected Response stream_id = "streamId-315624902" stream_token = b"122" expected_response = { "stream_id": stream_id, "stream_token": stream_token } expected_response = firestore_pb2.WriteResponse(**expected_response) # Mock the API response channel = ChannelStub(responses=[iter([expected_response])]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = firestore_client.FirestoreClient() # Setup Request database = client.database_root_path("[PROJECT]", "[DATABASE]") request = {"database": database} request = firestore_pb2.WriteRequest(**request) requests = [request] response = client.write(requests) resources = list(response) assert len(resources) == 1 assert expected_response == resources[0] assert len(channel.requests) == 1 actual_requests = channel.requests[0][1] assert len(actual_requests) == 1 actual_request = list(actual_requests)[0] assert request == actual_request
def test_write_exception(self): # Mock the API response channel = ChannelStub(responses=[CustomException()]) patch = mock.patch("google.api_core.grpc_helpers.create_channel") with patch as create_channel: create_channel.return_value = channel client = firestore_client.FirestoreClient() # Setup request database = client.database_root_path("[PROJECT]", "[DATABASE]") request = {"database": database} request = firestore_pb2.WriteRequest(**request) requests = [request] with pytest.raises(CustomException): client.write(requests)