Skip to content

Commit d257fce

Browse files
committed
Polish 629497d
* Fix tests in MongoItemWriterBuilderTests * Minor code refactoring in MongoItemWriter
1 parent 629497d commit d257fce

File tree

3 files changed

+71
-42
lines changed

3 files changed

+71
-42
lines changed

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java

Lines changed: 39 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,12 @@
5353
*
5454
* @author Michael Minella
5555
* @author Parikshit Dutta
56+
* @author Mahmoud Ben Hassine
5657
*
5758
*/
5859
public class MongoItemWriter<T> implements ItemWriter<T>, InitializingBean {
5960

61+
private static final String ID_KEY = "_id";
6062
private MongoOperations template;
6163
private final Object bufferKey;
6264
private String collection;
@@ -130,46 +132,50 @@ public void write(List<? extends T> items) throws Exception {
130132
* @param items the list of items to be persisted.
131133
*/
132134
protected void doWrite(List<? extends T> items) {
133-
if(! CollectionUtils.isEmpty(items)) {
134-
if(delete) {
135-
if(StringUtils.hasText(collection)) {
136-
for (Object object : items) {
137-
template.remove(object, collection);
138-
}
139-
}
140-
else {
141-
for (Object object : items) {
142-
template.remove(object);
143-
}
144-
}
135+
if (!CollectionUtils.isEmpty(items)) {
136+
if (this.delete) {
137+
delete(items);
145138
}
146139
else {
147-
BulkOperations bulkOperations = null;
148-
149-
if(StringUtils.hasText(collection)) {
150-
bulkOperations = template.bulkOps(BulkMode.ORDERED, collection);
151-
}
152-
else {
153-
bulkOperations = template.bulkOps(BulkMode.ORDERED, ClassUtils.getUserClass(items.get(0)));
154-
}
155-
156-
for (Object object : items) {
157-
Document document = new Document();
158-
159-
MongoConverter mongoConverter = template.getConverter();
160-
mongoConverter.write(object, document);
161-
162-
Query query = new Query();
163-
query.addCriteria(Criteria.where("_id").is((document.get("_id") != null)
164-
? document.get("_id") : new ObjectId()));
140+
saveOrUpdate(items);
141+
}
142+
}
143+
}
165144

166-
bulkOperations.replaceOne(query, document, new FindAndReplaceOptions().upsert());
167-
}
168-
bulkOperations.execute();
145+
private void delete(List<? extends T> items) {
146+
if (StringUtils.hasText(this.collection)) {
147+
for (Object item : items) {
148+
this.template.remove(item, this.collection);
149+
}
150+
}
151+
else {
152+
for (Object item : items) {
153+
this.template.remove(item);
169154
}
170155
}
171156
}
172157

158+
private void saveOrUpdate(List<? extends T> items) {
159+
BulkOperations bulkOperations;
160+
BulkMode bulkMode = BulkMode.ORDERED;
161+
if (StringUtils.hasText(this.collection)) {
162+
bulkOperations = this.template.bulkOps(bulkMode, this.collection);
163+
}
164+
else {
165+
bulkOperations = this.template.bulkOps(bulkMode, ClassUtils.getUserClass(items.get(0)));
166+
}
167+
MongoConverter mongoConverter = this.template.getConverter();
168+
FindAndReplaceOptions upsert = new FindAndReplaceOptions().upsert();
169+
for (Object item : items) {
170+
Document document = new Document();
171+
mongoConverter.write(item, document);
172+
Object objectId = document.get(ID_KEY) != null ? document.get(ID_KEY) : new ObjectId();
173+
Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId));
174+
bulkOperations.replaceOne(query, document, upsert);
175+
}
176+
bulkOperations.execute();
177+
}
178+
173179
private boolean transactionActive() {
174180
return TransactionSynchronizationManager.isActualTransactionActive();
175181
}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/MongoItemWriterTests.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import static org.junit.Assert.fail;
3939
import static org.mockito.ArgumentMatchers.any;
4040
import static org.mockito.ArgumentMatchers.anyString;
41+
import static org.mockito.ArgumentMatchers.eq;
4142
import static org.mockito.Mockito.when;
4243
import static org.mockito.Mockito.doAnswer;
4344
import static org.mockito.Mockito.mock;
@@ -48,6 +49,7 @@
4849
/**
4950
* @author Michael Minella
5051
* @author Parikshit Dutta
52+
* @author Mahmoud Ben Hassine
5153
*/
5254
@SuppressWarnings("serial")
5355
public class MongoItemWriterTests {
@@ -111,7 +113,7 @@ public void testWriteNoTransactionWithCollection() throws Exception {
111113

112114
writer.write(items);
113115

114-
verify(template).bulkOps(any(), anyString());
116+
verify(template).bulkOps(any(), eq("collection"));
115117
verify(bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any());
116118
}
117119

@@ -163,7 +165,7 @@ public void testWriteTransactionWithCollection() throws Exception {
163165
return null;
164166
});
165167

166-
verify(template).bulkOps(any(), anyString());
168+
verify(template).bulkOps(any(), eq("collection"));
167169
verify(bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any());
168170
}
169171

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilderTests.java

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2017 the original author or authors.
2+
* Copyright 2017-2020 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -19,31 +19,48 @@
1919
import java.util.Arrays;
2020
import java.util.List;
2121

22+
import org.bson.Document;
2223
import org.junit.Before;
2324
import org.junit.Test;
2425
import org.mockito.Mock;
2526
import org.mockito.MockitoAnnotations;
2627

2728
import org.springframework.batch.item.data.MongoItemWriter;
29+
import org.springframework.data.mongodb.core.BulkOperations;
2830
import org.springframework.data.mongodb.core.MongoOperations;
31+
import org.springframework.data.mongodb.core.convert.MongoConverter;
32+
import org.springframework.data.mongodb.core.query.Query;
2933

3034
import static org.junit.Assert.assertEquals;
3135
import static org.junit.Assert.fail;
36+
import static org.mockito.ArgumentMatchers.any;
37+
import static org.mockito.ArgumentMatchers.anyString;
38+
import static org.mockito.ArgumentMatchers.eq;
3239
import static org.mockito.Mockito.never;
40+
import static org.mockito.Mockito.times;
3341
import static org.mockito.Mockito.verify;
42+
import static org.mockito.Mockito.when;
3443

3544
/**
3645
* @author Glenn Renfro
46+
* @author Mahmoud Ben Hassine
3747
*/
3848
public class MongoItemWriterBuilderTests {
3949
@Mock
4050
private MongoOperations template;
51+
@Mock
52+
private BulkOperations bulkOperations;
53+
@Mock
54+
private MongoConverter mongoConverter;
4155

4256
private List<String> items;
4357

4458
@Before
4559
public void setUp() throws Exception {
4660
MockitoAnnotations.initMocks(this);
61+
when(this.template.bulkOps(any(), anyString())).thenReturn(this.bulkOperations);
62+
when(this.template.bulkOps(any(), any(Class.class))).thenReturn(this.bulkOperations);
63+
when(this.template.getConverter()).thenReturn(this.mongoConverter);
4764
this.items = Arrays.asList("foo", "bar");
4865
}
4966

@@ -52,8 +69,10 @@ public void testBasicWrite() throws Exception {
5269
MongoItemWriter<String> writer = new MongoItemWriterBuilder<String>().template(this.template).build();
5370
writer.write(this.items);
5471

55-
verify(this.template).save(this.items.get(0));
56-
verify(this.template).save(this.items.get(1));
72+
verify(this.template).bulkOps(any(), any(Class.class));
73+
verify(this.mongoConverter).write(eq(this.items.get(0)), any(Document.class));
74+
verify(this.mongoConverter).write(eq(this.items.get(1)), any(Document.class));
75+
verify(this.bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any());
5776
verify(this.template, never()).remove(this.items.get(0));
5877
verify(this.template, never()).remove(this.items.get(1));
5978
}
@@ -68,8 +87,8 @@ public void testDelete() throws Exception {
6887

6988
verify(this.template).remove(this.items.get(0));
7089
verify(this.template).remove(this.items.get(1));
71-
verify(this.template, never()).save(this.items.get(0));
72-
verify(this.template, never()).save(this.items.get(1));
90+
verify(this.template, never()).bulkOps(any(), any(Class.class));
91+
verify(this.mongoConverter, never()).write(any(), any());
7392
}
7493

7594
@Test
@@ -80,8 +99,10 @@ public void testWriteToCollection() throws Exception {
8099

81100
writer.write(this.items);
82101

83-
verify(this.template).save(this.items.get(0), "collection");
84-
verify(this.template).save(this.items.get(1), "collection");
102+
verify(this.template).bulkOps(any(), eq("collection"));
103+
verify(this.mongoConverter).write(eq(this.items.get(0)), any(Document.class));
104+
verify(this.mongoConverter).write(eq(this.items.get(1)), any(Document.class));
105+
verify(this.bulkOperations, times(2)).replaceOne(any(Query.class), any(Object.class), any());
85106
verify(this.template, never()).remove(this.items.get(0), "collection");
86107
verify(this.template, never()).remove(this.items.get(1), "collection");
87108
}

0 commit comments

Comments
 (0)