From 932871cb0e2584c1cb085b38e84656f387e2b976 Mon Sep 17 00:00:00 2001
From: Oliver Eilhard <oliver@eilhard.net>
Date: Fri, 7 Jan 2022 18:07:00 +0100
Subject: [PATCH] Allow multiple inner hits with CollapseBuilder

The `CollapseBuilder` now accepts multiple inner hits.

Close #1553
---
 search_collapse_builder.go      | 24 ++++++++++++++----------
 search_collapse_builder_test.go | 22 +++++++++++++++++++++-
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/search_collapse_builder.go b/search_collapse_builder.go
index 0de4eed61..b9488a566 100644
--- a/search_collapse_builder.go
+++ b/search_collapse_builder.go
@@ -9,7 +9,7 @@ package elastic
 // for details.
 type CollapseBuilder struct {
 	field                      string
-	innerHit                   *InnerHit
+	innerHits                  []*InnerHit
 	maxConcurrentGroupRequests *int
 }
 
@@ -25,8 +25,8 @@ func (b *CollapseBuilder) Field(field string) *CollapseBuilder {
 }
 
 // InnerHit option to expand the collapsed results.
-func (b *CollapseBuilder) InnerHit(innerHit *InnerHit) *CollapseBuilder {
-	b.innerHit = innerHit
+func (b *CollapseBuilder) InnerHit(innerHits ...*InnerHit) *CollapseBuilder {
+	b.innerHits = append(b.innerHits, innerHits...)
 	return b
 }
 
@@ -41,23 +41,27 @@ func (b *CollapseBuilder) MaxConcurrentGroupRequests(max int) *CollapseBuilder {
 func (b *CollapseBuilder) Source() (interface{}, error) {
 	// {
 	//   "field": "user",
-	//   "inner_hits": {
+	//   "inner_hits": [{
 	//     "name": "last_tweets",
 	//     "size": 5,
 	//     "sort": [{ "date": "asc" }]
-	//   },
+	//   }],
 	//   "max_concurrent_group_searches": 4
 	// }
 	src := map[string]interface{}{
 		"field": b.field,
 	}
 
-	if b.innerHit != nil {
-		hits, err := b.innerHit.Source()
-		if err != nil {
-			return nil, err
+	if len(b.innerHits) > 0 {
+		var innerHits []interface{}
+		for _, h := range b.innerHits {
+			hits, err := h.Source()
+			if err != nil {
+				return nil, err
+			}
+			innerHits = append(innerHits, hits)
 		}
-		src["inner_hits"] = hits
+		src["inner_hits"] = innerHits
 	}
 
 	if b.maxConcurrentGroupRequests != nil {
diff --git a/search_collapse_builder_test.go b/search_collapse_builder_test.go
index 0b74fadab..fc99b76d7 100644
--- a/search_collapse_builder_test.go
+++ b/search_collapse_builder_test.go
@@ -22,7 +22,27 @@ func TestCollapseBuilderSource(t *testing.T) {
 		t.Fatalf("marshaling to JSON failed: %v", err)
 	}
 	got := string(data)
-	expected := `{"field":"user","inner_hits":{"name":"last_tweets","size":5,"sort":[{"date":{"order":"asc"}}]},"max_concurrent_group_searches":4}`
+	expected := `{"field":"user","inner_hits":[{"name":"last_tweets","size":5,"sort":[{"date":{"order":"asc"}}]}],"max_concurrent_group_searches":4}`
+	if got != expected {
+		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+	}
+}
+
+func TestCollapseBuilderSourceMultipleInnerHits(t *testing.T) {
+	b := NewCollapseBuilder("user.id").
+		InnerHit(NewInnerHit().Name("largest_responses").Size(3).Sort("http.response.bytes", false)).
+		InnerHit(NewInnerHit().Name("most_recent").Size(4).Sort("@timestamp", false)).
+		MaxConcurrentGroupRequests(5)
+	src, err := b.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"field":"user.id","inner_hits":[{"name":"largest_responses","size":3,"sort":[{"http.response.bytes":{"order":"desc"}}]},{"name":"most_recent","size":4,"sort":[{"@timestamp":{"order":"desc"}}]}],"max_concurrent_group_searches":5}`
 	if got != expected {
 		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
 	}