diff --git a/machine_test.go b/machine_test.go index cf0d12e7..dfe84d33 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 25de5846..baa90d0c 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)