From f6a7564ec870a57c1b22c101578dd1edd9ba43b7 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 4 Oct 2021 17:40:21 +0000
Subject: [PATCH] Add more test cases to prove that peers and shared peers work
 properly

---
 machine_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 sharing_test.go |  9 +++++++++
 2 files changed, 50 insertions(+)

diff --git a/machine_test.go b/machine_test.go
index cf0d12e..dfe84d3 100644
--- a/machine_test.go
+++ b/machine_test.go
@@ -2,6 +2,7 @@ package headscale
 
 import (
 	"encoding/json"
+	"strconv"
 
 	"gopkg.in/check.v1"
 )
@@ -116,3 +117,43 @@ func (s *Suite) TestHardDeleteMachine(c *check.C) {
 	_, err = h.GetMachine(n.Name, "testmachine3")
 	c.Assert(err, check.NotNil)
 }
+
+func (s *Suite) TestGetDirectPeers(c *check.C) {
+	n, err := h.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
+	c.Assert(err, check.IsNil)
+
+	_, err = h.GetMachineByID(0)
+	c.Assert(err, check.NotNil)
+
+	for i := 0; i <= 10; i++ {
+		m := Machine{
+			ID:             uint64(i),
+			MachineKey:     "foo" + strconv.Itoa(i),
+			NodeKey:        "bar" + strconv.Itoa(i),
+			DiscoKey:       "faa" + strconv.Itoa(i),
+			Name:           "testmachine" + strconv.Itoa(i),
+			NamespaceID:    n.ID,
+			Registered:     true,
+			RegisterMethod: "authKey",
+			AuthKeyID:      uint(pak.ID),
+		}
+		h.db.Save(&m)
+	}
+
+	m1, err := h.GetMachineByID(0)
+	c.Assert(err, check.IsNil)
+
+	_, err = m1.GetHostInfo()
+	c.Assert(err, check.IsNil)
+
+	peers, err := h.getDirectPeers(m1)
+	c.Assert(err, check.IsNil)
+
+	c.Assert(len(peers), check.Equals, 9)
+	c.Assert(peers[0].Name, check.Equals, "testmachine2")
+	c.Assert(peers[5].Name, check.Equals, "testmachine7")
+	c.Assert(peers[8].Name, check.Equals, "testmachine10")
+}
diff --git a/sharing_test.go b/sharing_test.go
index 25de584..baa90d0 100644
--- a/sharing_test.go
+++ b/sharing_test.go
@@ -245,6 +245,7 @@ func (s *Suite) TestDoNotIncludeRoutesOnShared(c *check.C) {
 	p1sAfter, err := h.getPeers(m1)
 	c.Assert(err, check.IsNil)
 	c.Assert(len(p1sAfter), check.Equals, 1)
+	c.Assert(p1sAfter[0].Name, check.Equals, "test_get_shared_nodes_2")
 }
 
 func (s *Suite) TestComplexSharingAcrossNamespaces(c *check.C) {
@@ -343,6 +344,7 @@ func (s *Suite) TestComplexSharingAcrossNamespaces(c *check.C) {
 	p1s, err := h.getPeers(m1)
 	c.Assert(err, check.IsNil)
 	c.Assert(len(p1s), check.Equals, 1) // nodes 1 and 4
+	c.Assert(p1s[0].Name, check.Equals, "test_get_shared_nodes_4")
 
 	err = h.AddSharedMachineToNamespace(m2, n1)
 	c.Assert(err, check.IsNil)
@@ -350,6 +352,13 @@ func (s *Suite) TestComplexSharingAcrossNamespaces(c *check.C) {
 	p1sAfter, err := h.getPeers(m1)
 	c.Assert(err, check.IsNil)
 	c.Assert(len(p1sAfter), check.Equals, 2) // nodes 1, 2, 4
+	c.Assert(p1sAfter[0].Name, check.Equals, "test_get_shared_nodes_2")
+	c.Assert(p1sAfter[1].Name, check.Equals, "test_get_shared_nodes_4")
+
+	node1shared, err := h.getShared(m1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(node1shared), check.Equals, 1) // nodes 1, 2, 4
+	c.Assert(node1shared[0].Name, check.Equals, "test_get_shared_nodes_2")
 
 	pAlone, err := h.getPeers(m3)
 	c.Assert(err, check.IsNil)