From 81491843f108a6f0eb15ab672778092675ec14a4 Mon Sep 17 00:00:00 2001
From: Colin Goodheart-Smithe <colings86@users.noreply.github.com>
Date: Mon, 5 Oct 2015 14:10:16 +0100
Subject: [PATCH] test fixes

---
 .../elasticsearch/action/count/CountRequest.java | 16 ++++++++--------
 .../action/count/CountRequestBuilder.java        |  6 +-----
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/action/count/CountRequest.java b/core/src/main/java/org/elasticsearch/action/count/CountRequest.java
index a74ac54e6ce..4e03b03ee10 100644
--- a/core/src/main/java/org/elasticsearch/action/count/CountRequest.java
+++ b/core/src/main/java/org/elasticsearch/action/count/CountRequest.java
@@ -58,7 +58,7 @@ public class CountRequest extends BroadcastRequest<CountRequest> {
 
     private int terminateAfter = DEFAULT_TERMINATE_AFTER;
 
-    private QueryBuilder<?> queryBuilder = null;
+    private SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
 
     /**
      * Constructs a new count request against the provided indices. No indices provided means it will
@@ -88,13 +88,11 @@ public class CountRequest extends BroadcastRequest<CountRequest> {
     /**
      * The query to execute
      */
-    public CountRequest query(QueryBuilder<?> queryBuilder) {
-        this.queryBuilder = queryBuilder;
+    public CountRequest query(QueryBuilder queryBuilder) {
+        this.searchSourceBuilder = new SearchSourceBuilder().query(queryBuilder);
         return this;
     }
 
-
-
     /**
      * The types of documents the query will run against. Defaults to all types.
      */
@@ -171,16 +169,18 @@ public class CountRequest extends BroadcastRequest<CountRequest> {
     public String toString() {
         String sSource = "_na_";
         try {
-            sSource = XContentHelper.toString(queryBuilder);
+            sSource = XContentHelper.toString(searchSourceBuilder);
         } catch (Exception e) {
             // ignore
         }
         return "[" + Arrays.toString(indices) + "]" + Arrays.toString(types) + ", source[" + sSource + "]";
     }
 
+    public String sourceBuilderString() {
+        return searchSourceBuilder.toString();
+    }
+
     public SearchRequest toSearchRequest() {
-        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
-        searchSourceBuilder.query(queryBuilder);
         searchSourceBuilder.size(0);
         if (minScore() != DEFAULT_MIN_SCORE) {
             searchSourceBuilder.minScore(minScore());
diff --git a/core/src/main/java/org/elasticsearch/action/count/CountRequestBuilder.java b/core/src/main/java/org/elasticsearch/action/count/CountRequestBuilder.java
index c10c54e14de..f11996f6b27 100644
--- a/core/src/main/java/org/elasticsearch/action/count/CountRequestBuilder.java
+++ b/core/src/main/java/org/elasticsearch/action/count/CountRequestBuilder.java
@@ -23,7 +23,6 @@ import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.search.SearchAction;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.support.DelegatingActionListener;
-import org.elasticsearch.action.support.QuerySourceBuilder;
 import org.elasticsearch.action.support.broadcast.BroadcastOperationRequestBuilder;
 import org.elasticsearch.client.ElasticsearchClient;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -107,9 +106,6 @@ public class CountRequestBuilder extends BroadcastOperationRequestBuilder<CountR
 
     @Override
     public String toString() {
-        if (request != null) {
-            return request.toString();
-        }
-        return new QuerySourceBuilder().toString();
+        return request.sourceBuilderString();
     }
 }