From aa5bdcd3faf11ac674ce17ab542d8c7d5d31ea50 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@pivotal.io>
Date: Thu, 3 Mar 2016 17:34:33 +0200
Subject: [PATCH] DATAMONGO-1387 - Polishing.

Added a few more tests and append values if present on Query.

Original Pull Request: #345
---
 .../data/mongodb/core/query/BasicQuery.java   | 17 +++++--
 .../core/query/BasicQueryUnitTests.java       | 51 +++++++++++++++----
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/BasicQuery.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/BasicQuery.java
index f3e37d0a4..a871162fc 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/BasicQuery.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/BasicQuery.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2014 the original author or authors.
+ * Copyright 2010-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -71,11 +71,20 @@ public class BasicQuery extends Query {
 
 	@Override
 	public DBObject getFieldsObject() {
-		if(fieldsObject != null) {
-			return fieldsObject;
-		} else {
+
+		if (fieldsObject == null) {
 			return super.getFieldsObject();
 		}
+
+		if (super.getFieldsObject() != null) {
+
+			DBObject combinedFieldsObject = new BasicDBObject();
+			combinedFieldsObject.putAll(fieldsObject);
+			combinedFieldsObject.putAll(super.getFieldsObject());
+			return combinedFieldsObject;
+		}
+
+		return fieldsObject;
 	}
 
 	@Override
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/BasicQueryUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/BasicQueryUnitTests.java
index 5c69e9851..e933768b1 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/BasicQueryUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/BasicQueryUnitTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2011-2014 the original author or authors.
+ * Copyright 2011-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,8 +18,7 @@ package org.springframework.data.mongodb.core.query;
 import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
-import nl.jqno.equalsverifier.EqualsVerifier;
-import nl.jqno.equalsverifier.Warning;
+import static org.springframework.data.mongodb.test.util.IsBsonObject.*;
 
 import org.junit.Test;
 import org.springframework.data.domain.Sort.Direction;
@@ -27,6 +26,9 @@ import org.springframework.data.domain.Sort.Direction;
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
 
+import nl.jqno.equalsverifier.EqualsVerifier;
+import nl.jqno.equalsverifier.Warning;
+
 /**
  * Unit tests for {@link BasicQuery}.
  * 
@@ -138,21 +140,48 @@ public class BasicQueryUnitTests {
 		assertThat(query1, is(not(equalTo(query2))));
 		assertThat(query1.hashCode(), is(not(query2.hashCode())));
 	}
-	
+
+	/**
+	 * @see DATAMONGO-1387
+	 */
+	@Test
+	public void returnsFieldsCorrectly() {
+
+		String qry = "{ \"name\" : \"Thomas\"}";
+		String fields = "{\"name\":1, \"age\":1}";
+
+		BasicQuery query1 = new BasicQuery(qry, fields);
+
+		assertThat(query1.getFieldsObject(), isBsonObject().containing("name").containing("age"));
+	}
+
 	/**
 	 * @see DATAMONGO-1387
 	 */
 	@Test
 	public void handlesFieldsIncludeCorrectly() {
-		
+
 		String qry = "{ \"name\" : \"Thomas\"}";
-		
+
 		BasicQuery query1 = new BasicQuery(qry);
 		query1.fields().include("name");
-		
-		DBObject fieldsObject = query1.getFieldsObject();
-		fieldsObject.containsField("name");
-		assertThat(query1.getFieldsObject(), notNullValue());
-		assertThat(query1.getFieldsObject().containsField("name"), is(true));
+
+		assertThat(query1.getFieldsObject(), isBsonObject().containing("name"));
 	}
+
+	/**
+	 * @see DATAMONGO-1387
+	 */
+	@Test
+	public void combinesFieldsIncludeCorrectly() {
+
+		String qry = "{ \"name\" : \"Thomas\"}";
+		String fields = "{\"name\":1, \"age\":1}";
+
+		BasicQuery query1 = new BasicQuery(qry, fields);
+		query1.fields().include("gender");
+
+		assertThat(query1.getFieldsObject(), isBsonObject().containing("name").containing("age").containing("gender"));
+	}
+
 }