Skip to content

Commit 7e2ca54

Browse files
committed
test(batch): add unit tests for TaskHandle
1 parent 7e6d6c6 commit 7e2ca54

1 file changed

Lines changed: 164 additions & 0 deletions

File tree

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
package io.weaviate.client6.v1.api.collections.batch;
2+
3+
import java.util.Map;
4+
import java.util.Optional;
5+
import java.util.concurrent.CompletableFuture;
6+
7+
import org.assertj.core.api.AbstractObjectAssert;
8+
import org.assertj.core.api.Assertions;
9+
import org.assertj.core.api.InstanceOfAssertFactories;
10+
import org.junit.Test;
11+
12+
import com.google.protobuf.GeneratedMessage;
13+
import com.google.protobuf.GeneratedMessageV3;
14+
15+
import io.weaviate.client6.v1.api.collections.CollectionHandleDefaults;
16+
import io.weaviate.client6.v1.api.collections.WeaviateObject;
17+
import io.weaviate.client6.v1.api.collections.data.BatchReference;
18+
import io.weaviate.client6.v1.api.collections.data.InsertManyRequest;
19+
import io.weaviate.client6.v1.api.collections.data.ObjectReference;
20+
import io.weaviate.client6.v1.internal.orm.CollectionDescriptor;
21+
22+
public class TaskHandleTest {
23+
private static final WeaviateObject<Map<String, Object>> OBJECT = WeaviateObject.of();
24+
private static final BatchReference REFERENCE = new BatchReference(
25+
"Songs", "hasAwards", "song-1",
26+
ObjectReference.collection("GrammyAwards", "grammy-1"));
27+
28+
private static final GeneratedMessage.ExtendableMessage<GeneratedMessageV3> REFERENCE_PROTO = InsertManyRequest
29+
.buildReference(REFERENCE, Optional.empty());
30+
private static final GeneratedMessage.ExtendableMessage<GeneratedMessageV3> OBJECT_PROTO = InsertManyRequest
31+
.buildObject(OBJECT,
32+
CollectionDescriptor.ofMap("Songs"),
33+
CollectionHandleDefaults.of(CollectionHandleDefaults.none()));
34+
35+
@Test
36+
public void testTaskHandle_WeaviateObject_success() {
37+
TaskHandle taskHandle = new TaskHandle(OBJECT, OBJECT_PROTO);
38+
39+
Assertions.assertThat(taskHandle)
40+
.returns(OBJECT.uuid(), TaskHandle::id)
41+
.returns(0, TaskHandle::timesRetried);
42+
43+
assertAcked(taskHandle, false);
44+
assertHasResult(taskHandle, false);
45+
46+
taskHandle.setAcked();
47+
48+
assertAcked(taskHandle, true);
49+
assertHasResult(taskHandle, false);
50+
51+
taskHandle.setSuccess();
52+
53+
assertAcked(taskHandle, true);
54+
assertHasResult(taskHandle, true)
55+
.extracting(future -> future.getNow(null)).isNotNull()
56+
.extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
57+
.isEmpty();
58+
}
59+
60+
@Test
61+
public void testTaskHandle_WeaviateObject_error() {
62+
TaskHandle taskHandle = new TaskHandle(OBJECT, OBJECT_PROTO);
63+
64+
Assertions.assertThat(taskHandle)
65+
.returns(OBJECT.uuid(), TaskHandle::id)
66+
.returns(0, TaskHandle::timesRetried);
67+
68+
assertAcked(taskHandle, false);
69+
assertHasResult(taskHandle, false);
70+
71+
taskHandle.setAcked();
72+
73+
assertAcked(taskHandle, true);
74+
assertHasResult(taskHandle, false);
75+
76+
taskHandle.setError("Whaam!");
77+
78+
assertAcked(taskHandle, true);
79+
assertHasResult(taskHandle, true)
80+
.extracting(future -> future.getNow(null)).isNotNull()
81+
.extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
82+
.get().isEqualTo("Whaam!");
83+
}
84+
85+
@Test
86+
public void testTaskHandle_BatchReference_success() {
87+
TaskHandle taskHandle = new TaskHandle(REFERENCE, REFERENCE_PROTO);
88+
89+
Assertions.assertThat(taskHandle)
90+
.returns(REFERENCE.target().beacon(), TaskHandle::id)
91+
.returns(0, TaskHandle::timesRetried);
92+
93+
assertAcked(taskHandle, false);
94+
assertHasResult(taskHandle, false);
95+
96+
taskHandle.setAcked();
97+
98+
assertAcked(taskHandle, true);
99+
assertHasResult(taskHandle, false);
100+
101+
taskHandle.setSuccess();
102+
103+
assertAcked(taskHandle, true);
104+
assertHasResult(taskHandle, true)
105+
.extracting(future -> future.getNow(null)).isNotNull()
106+
.extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
107+
.isEmpty();
108+
}
109+
110+
@Test
111+
public void testTaskHandle_BatchReference_error() {
112+
TaskHandle taskHandle = new TaskHandle(REFERENCE, REFERENCE_PROTO);
113+
114+
Assertions.assertThat(taskHandle)
115+
.returns(REFERENCE.target().beacon(), TaskHandle::id)
116+
.returns(0, TaskHandle::timesRetried);
117+
118+
assertAcked(taskHandle, false);
119+
assertHasResult(taskHandle, false);
120+
121+
taskHandle.setAcked();
122+
123+
assertAcked(taskHandle, true);
124+
assertHasResult(taskHandle, false);
125+
126+
taskHandle.setError("Whaam!");
127+
128+
assertAcked(taskHandle, true);
129+
assertHasResult(taskHandle, true)
130+
.extracting(future -> future.getNow(null)).isNotNull()
131+
.extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
132+
.get().isEqualTo("Whaam!");
133+
}
134+
135+
@Test
136+
public void testTaskHandle_retry() {
137+
TaskHandle taskHandle = new TaskHandle(OBJECT, OBJECT_PROTO);
138+
Assertions.assertThat(taskHandle).returns(0, TaskHandle::timesRetried);
139+
140+
TaskHandle retried;
141+
Assertions.assertThat(retried = taskHandle.retry())
142+
.returns(1, TaskHandle::timesRetried)
143+
.extracting(TaskHandle::data).isEqualTo(taskHandle.data());
144+
145+
Assertions.assertThat(retried = retried.retry())
146+
.returns(2, TaskHandle::timesRetried)
147+
.extracting(TaskHandle::data).isEqualTo(taskHandle.data());
148+
}
149+
150+
private void assertAcked(TaskHandle taskHandle, boolean expect) {
151+
Assertions.assertThat(taskHandle)
152+
.extracting(TaskHandle::isAcked)
153+
.as("expect is acked")
154+
.returns(expect, CompletableFuture::isDone);
155+
}
156+
157+
private AbstractObjectAssert<?, CompletableFuture<TaskHandle.Result>> assertHasResult(TaskHandle taskHandle,
158+
boolean expect) {
159+
return Assertions.assertThat(taskHandle)
160+
.extracting(TaskHandle::result)
161+
.as("expect has result")
162+
.returns(expect, CompletableFuture::isDone);
163+
}
164+
}

0 commit comments

Comments
 (0)