From ca0304f527d203e4032e70b38fc4f7299f9b64a7 Mon Sep 17 00:00:00 2001
From: Sam Woodard <samwoodard@nestlabs.com>
Date: Sun, 12 Jul 2015 07:42:52 -0700
Subject: [PATCH] bug fix: don't forget some included nodes

---
 request_test.go  | 43 +++++++++++++++++++++++++++++++++++++++++++
 response.go      | 32 ++++++++++++++++++++------------
 response_test.go | 11 ++++++-----
 3 files changed, 69 insertions(+), 17 deletions(-)

diff --git a/request_test.go b/request_test.go
index c63bb63..4143f01 100644
--- a/request_test.go
+++ b/request_test.go
@@ -93,6 +93,26 @@ func TestUnmarshalNestedRelationships(t *testing.T) {
 	}
 }
 
+func TestUnmarshalRelationshipsSerializedEmbedded(t *testing.T) {
+	out := sampleSerializedEmbeddedTestModel()
+
+	if out.CurrentPost == nil {
+		t.Fatalf("Current post was not materialized")
+	}
+
+	if out.CurrentPost.Title != "Foo" || out.CurrentPost.Body != "Bar" {
+		t.Fatalf("Attributes were not set")
+	}
+
+	if len(out.Posts) != 2 {
+		t.Fatalf("There should have been 2 posts")
+	}
+
+	if out.Posts[0].LatestComment.Body != "foo" {
+		t.Fatalf("The comment body was not set")
+	}
+}
+
 func TestUnmarshalNestedRelationshipsEmbedded(t *testing.T) {
 	out := bytes.NewBuffer(nil)
 	if err := MarshalOnePayloadEmbedded(out, testModel()); err != nil {
@@ -284,6 +304,10 @@ func testModel() *Blog {
 						Body: "bar",
 					},
 				},
+				LatestComment: &Comment{
+					Id:   1,
+					Body: "foo",
+				},
 			},
 			&Post{
 				Id:    2,
@@ -299,6 +323,10 @@ func testModel() *Blog {
 						Body: "bas",
 					},
 				},
+				LatestComment: &Comment{
+					Id:   1,
+					Body: "foo",
+				},
 			},
 		},
 		CurrentPost: &Post{
@@ -315,6 +343,10 @@ func testModel() *Blog {
 					Body: "bar",
 				},
 			},
+			LatestComment: &Comment{
+				Id:   1,
+				Body: "foo",
+			},
 		},
 	}
 }
@@ -327,3 +359,14 @@ func samplePayloadWithSideloaded() io.Reader {
 
 	return out
 }
+
+func sampleSerializedEmbeddedTestModel() *Blog {
+	out := bytes.NewBuffer(nil)
+	MarshalOnePayloadEmbedded(out, testModel())
+
+	blog := new(Blog)
+
+	UnmarshalPayload(out, blog)
+
+	return blog
+}
diff --git a/response.go b/response.go
index 5fb7652..762b2ab 100644
--- a/response.go
+++ b/response.go
@@ -161,12 +161,12 @@ func visitModelNode(model interface{}, sideload bool) (*Node, []*Node, error) {
 				}
 
 				if isSlice {
-					relationship, err := visitModelNodeRelationships(args[1], fieldValue, sideload)
-					d := relationship[args[1]].Data
+					relationship, incl, err := visitModelNodeRelationships(args[1], fieldValue, sideload)
+					d := relationship.Data
 
 					if err == nil {
 						if sideload {
-							included = append(included, d...)
+							included = append(included, incl...)
 							shallowNodes := make([]*Node, 0)
 							for _, node := range d {
 								shallowNodes = append(shallowNodes, toShallowNode(node))
@@ -174,17 +174,18 @@ func visitModelNode(model interface{}, sideload bool) (*Node, []*Node, error) {
 
 							node.Relationships[args[1]] = &RelationshipManyNode{Data: shallowNodes}
 						} else {
-							node.Relationships[args[1]] = &RelationshipManyNode{Data: d}
+							node.Relationships[args[1]] = relationship
 						}
 					} else {
 						er = err
 						return false
 					}
 				} else {
-					relationship, inc, err := visitModelNode(fieldValue.Interface(), sideload)
+					relationship, incl, err := visitModelNode(fieldValue.Interface(), sideload)
 					if err == nil {
 						if sideload {
-							included = append(included, inc...)
+							included = append(included, incl...)
+							included = append(included, relationship)
 							node.Relationships[args[1]] = &RelationshipOneNode{Data: toShallowNode(relationship)}
 						} else {
 							node.Relationships[args[1]] = &RelationshipOneNode{Data: relationship}
@@ -218,22 +219,29 @@ func toShallowNode(node *Node) *Node {
 	}
 }
 
-func visitModelNodeRelationships(relationName string, models reflect.Value, sideload bool) (map[string]*RelationshipManyNode, error) {
-	m := make(map[string]*RelationshipManyNode)
+func visitModelNodeRelationships(relationName string, models reflect.Value, sideload bool) (*RelationshipManyNode, []*Node, error) {
 	nodes := make([]*Node, 0)
 
+	var included []*Node
+	if sideload {
+		included = make([]*Node, 0)
+	}
+
 	for i := 0; i < models.Len(); i++ {
-		node, _, err := visitModelNode(models.Index(i).Interface(), sideload)
+		node, incl, err := visitModelNode(models.Index(i).Interface(), sideload)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 
 		nodes = append(nodes, node)
+		included = append(included, incl...)
 	}
 
-	m[relationName] = &RelationshipManyNode{Data: nodes}
+	included = append(included, nodes...)
 
-	return m, nil
+	n := &RelationshipManyNode{Data: nodes}
+
+	return n, included, nil
 }
 
 func deleteNode(a []*Node, i int) []*Node {
diff --git a/response_test.go b/response_test.go
index ba3a4c7..c48721d 100644
--- a/response_test.go
+++ b/response_test.go
@@ -19,11 +19,12 @@ type Blog struct {
 }
 
 type Post struct {
-	Id       int        `jsonapi:"primary,posts"`
-	BlogId   int        `jsonapi:"attr,blog_id"`
-	Title    string     `jsonapi:"attr,title"`
-	Body     string     `jsonapi:"attr,body"`
-	Comments []*Comment `jsonapi:"relation,comments"`
+	Id            int        `jsonapi:"primary,posts"`
+	BlogId        int        `jsonapi:"attr,blog_id"`
+	Title         string     `jsonapi:"attr,title"`
+	Body          string     `jsonapi:"attr,body"`
+	Comments      []*Comment `jsonapi:"relation,comments"`
+	LatestComment *Comment   `jsonapi:"relation,latest_comment"`
 }
 
 type Comment struct {