From 070934743f88d0a209a1290418518864e062988b Mon Sep 17 00:00:00 2001
From: Wichert Akkerman <wichert@wiggy.net>
Date: Tue, 29 Oct 2019 10:34:33 +0100
Subject: [PATCH] Add tests for iq result routes

---
 router_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/router_test.go b/router_test.go
index 98a4697..138999c 100644
--- a/router_test.go
+++ b/router_test.go
@@ -2,7 +2,9 @@ package xmpp
 
 import (
 	"bytes"
+	"context"
 	"encoding/xml"
+	"runtime"
 	"testing"
 
 	"gosrc.io/xmpp/stanza"
@@ -11,6 +13,58 @@ import (
 // ============================================================================
 // Test route & matchers
 
+func TestIqResultRoutes(t *testing.T) {
+	t.Parallel()
+	router := NewRouter()
+
+	if router.iqResultRoutes == nil {
+		t.Fatal("NewRouter does not initialize isResultRoutes")
+	}
+
+	// Check other IQ does not matcah
+	conn := NewSenderMock()
+	iq := stanza.NewIQ(stanza.Attrs{Type: stanza.IQTypeResult, Id: "4321"})
+	router.NewIqResultRoute(context.Background(), "1234").HandlerFunc(func(s Sender, p stanza.Packet) {
+		_ = s.SendRaw(successFlag)
+	})
+	if conn.String() == successFlag {
+		t.Fatal("IQ result with wrong ID was matched")
+	}
+
+	// Check if the IQ handler was called
+	conn = NewSenderMock()
+	iq = stanza.NewIQ(stanza.Attrs{Type: stanza.IQTypeResult, Id: "1234"})
+	router.route(conn, iq)
+	if conn.String() != successFlag {
+		t.Fatal("IQ result was not matched")
+	}
+
+	// The match must only happen once, so we if receive the same package again it
+	// must not be matched.
+	conn = NewSenderMock()
+	router.route(conn, iq)
+	if conn.String() == successFlag {
+		t.Fatal("IQ result was matched twice")
+	}
+
+	// After cancelling a route it should no longer match
+	conn = NewSenderMock()
+	ctx, cancel := context.WithCancel(context.Background())
+	iq = stanza.NewIQ(stanza.Attrs{Type: stanza.IQTypeResult, Id: "1234"})
+	router.NewIqResultRoute(ctx, "1234").HandlerFunc(func(s Sender, p stanza.Packet) {
+		_ = s.SendRaw(successFlag)
+	}).TimeoutHandlerFunc(func(err error) {
+		conn.SendRaw(cancelledFlag)
+	})
+	cancel()
+	// Yield the processor so the cancellation goroutine is triggered
+	runtime.Gosched()
+	router.route(conn, iq)
+	if conn.String() != cancelledFlag {
+		t.Fatal("IQ result route was matched after cancellation")
+	}
+}
+
 func TestNameMatcher(t *testing.T) {
 	router := NewRouter()
 	router.HandleFunc("message", func(s Sender, p stanza.Packet) {
@@ -211,7 +265,8 @@ func TestCatchallMatcher(t *testing.T) {
 // ============================================================================
 // SenderMock
 
-var successFlag = "matched"
+const successFlag = "matched"
+const cancelledFlag = "cancelled"
 
 type SenderMock struct {
 	buffer *bytes.Buffer