From 026e5e6fe198efef7c8a693f874d0f422731d580 Mon Sep 17 00:00:00 2001
From: Mickael Remond <mremond@process-one.net>
Date: Tue, 9 Jul 2019 16:59:54 +0200
Subject: [PATCH] Add helpers for IQ DiscoItems

---
 stanza/iq_disco.go      | 33 +++++++++++++++++++++++++--
 stanza/iq_disco_test.go | 50 ++++++++++++++++++++++++++++++++---------
 stanza/xmpp_test.go     | 25 +++++++++++++++++++++
 3 files changed, 95 insertions(+), 13 deletions(-)

diff --git a/stanza/iq_disco.go b/stanza/iq_disco.go
index 0097c05..cc94756 100644
--- a/stanza/iq_disco.go
+++ b/stanza/iq_disco.go
@@ -56,8 +56,9 @@ func (d *DiscoInfo) AddFeatures(namespace ...string) {
 	}
 }
 
-func (d *DiscoInfo) SetNode(node string) {
+func (d *DiscoInfo) SetNode(node string) *DiscoInfo {
 	d.Node = node
+	return d
 }
 
 func (d *DiscoInfo) SetIdentities(ident ...Identity) *DiscoInfo {
@@ -66,6 +67,7 @@ func (d *DiscoInfo) SetIdentities(ident ...Identity) *DiscoInfo {
 }
 
 func (d *DiscoInfo) SetFeatures(namespace ...string) *DiscoInfo {
+	d.Features = []Feature{}
 	for _, ns := range namespace {
 		d.Features = append(d.Features, Feature{Var: ns})
 	}
@@ -104,11 +106,38 @@ func (d *DiscoItems) Namespace() string {
 	return d.XMLName.Space
 }
 
+// ---------------
+// Builder helpers
+
+// DiscoItems builds a default DiscoItems payload
+func (iq *IQ) DiscoItems() *DiscoItems {
+	d := DiscoItems{
+		XMLName: xml.Name{Space: "http://jabber.org/protocol/disco#items", Local: "query"},
+	}
+	iq.Payload = &d
+	return &d
+}
+
+func (d *DiscoItems) SetNode(node string) *DiscoItems {
+	d.Node = node
+	return d
+}
+
+func (d *DiscoItems) AddItem(jid, node, name string) *DiscoItems {
+	item := DiscoItem{
+		JID:  jid,
+		Node: node,
+		Name: name,
+	}
+	d.Items = append(d.Items, item)
+	return d
+}
+
 type DiscoItem struct {
 	XMLName xml.Name `xml:"item"`
-	Name    string   `xml:"name,attr,omitempty"`
 	JID     string   `xml:"jid,attr,omitempty"`
 	Node    string   `xml:"node,attr,omitempty"`
+	Name    string   `xml:"name,attr,omitempty"`
 }
 
 // ============================================================================
diff --git a/stanza/iq_disco_test.go b/stanza/iq_disco_test.go
index 8b76767..47e2653 100644
--- a/stanza/iq_disco_test.go
+++ b/stanza/iq_disco_test.go
@@ -7,29 +7,21 @@ import (
 	"gosrc.io/xmpp/stanza"
 )
 
-func TestDiscoInfoBuilder(t *testing.T) {
+func TestDiscoInfo_Builder(t *testing.T) {
 	iq := stanza.NewIQ(stanza.Attrs{Type: "get", To: "service.localhost", Id: "disco-get-1"})
 	disco := iq.DiscoInfo()
 	disco.AddIdentity("Test Component", "gateway", "service")
 	disco.AddFeatures(stanza.NSDiscoInfo, stanza.NSDiscoItems, "jabber:iq:version", "urn:xmpp:delegation:1")
 
-	// Marshall
-	data, err := xml.Marshal(iq)
+	parsedIQ, err := marshallUnmarshall(t, iq)
 	if err != nil {
-		t.Errorf("cannot marshal xml structure: %s", err)
 		return
 	}
 
-	// Unmarshall
-	var parsedIQ stanza.IQ
-	if err = xml.Unmarshal(data, &parsedIQ); err != nil {
-		t.Errorf("Unmarshal(%s) returned error: %s", data, err)
-	}
-
 	// Check result
 	pp, ok := parsedIQ.Payload.(*stanza.DiscoInfo)
 	if !ok {
-		t.Errorf("Parsed stanza does not contain an IQ payload")
+		t.Errorf("Parsed stanza does not contain correct IQ payload")
 	}
 
 	// Check features
@@ -53,3 +45,39 @@ func TestDiscoInfoBuilder(t *testing.T) {
 		}
 	}
 }
+
+// Implements XEP-0030 example 17
+// https://xmpp.org/extensions/xep-0030.html#example-17
+func TestDiscoItems_Builder(t *testing.T) {
+	iq := stanza.NewIQ(stanza.Attrs{Type: "result", From: "catalog.shakespeare.lit",
+		To: "romeo@montague.net/orchard", Id: "items-2"})
+	iq.DiscoItems().
+		AddItem("catalog.shakespeare.lit", "books", "Books by and about Shakespeare").
+		AddItem("catalog.shakespeare.lit", "clothing", "Wear your literary taste with pride").
+		AddItem("catalog.shakespeare.lit", "music", "Music from the time of Shakespeare")
+
+	parsedIQ, err := marshallUnmarshall(t, iq)
+	if err != nil {
+		return
+	}
+
+	// Check result
+	pp, ok := parsedIQ.Payload.(*stanza.DiscoItems)
+	if !ok {
+		t.Errorf("Parsed stanza does not contain correct IQ payload")
+	}
+
+	// Check items
+	items := []stanza.DiscoItem{{xml.Name{}, "catalog.shakespeare.lit", "books", "Books by and about Shakespeare"},
+		{xml.Name{}, "catalog.shakespeare.lit", "clothing", "Wear your literary taste with pride"},
+		{xml.Name{}, "catalog.shakespeare.lit", "music", "Music from the time of Shakespeare"}}
+	if len(pp.Items) != len(items) {
+		t.Errorf("Items length mismatch: %#v", pp.Items)
+	} else {
+		for i, item := range pp.Items {
+			if item.JID != items[i].JID {
+				t.Errorf("JID Mismatch (expected: %s): %s", items[i].JID, item.JID)
+			}
+		}
+	}
+}
diff --git a/stanza/xmpp_test.go b/stanza/xmpp_test.go
index 611948a..7295b7b 100644
--- a/stanza/xmpp_test.go
+++ b/stanza/xmpp_test.go
@@ -2,10 +2,35 @@ package stanza_test
 
 import (
 	"encoding/xml"
+	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"gosrc.io/xmpp/stanza"
 )
 
+// ============================================================================
+// Marshaller / unmarshaller test
+
+func marshallUnmarshall(t *testing.T, iq stanza.IQ) (*stanza.IQ, error) {
+	// Marshall
+	data, err := xml.Marshal(iq)
+	if err != nil {
+		t.Errorf("cannot marshal iq: %s\n%#v", err, iq)
+		return nil, err
+	}
+
+	// Unmarshall
+	var parsedIQ stanza.IQ
+	err = xml.Unmarshal(data, &parsedIQ)
+	if err != nil {
+		t.Errorf("Unmarshal returned error: %s\n%s", err, data)
+	}
+	return &parsedIQ, err
+}
+
+// ============================================================================
+// XML structs comparison
+
 // Compare iq structure but ignore empty namespace as they are set properly on
 // marshal / unmarshal. There is no need to manage them on the manually
 // crafted structure.