Factor out some commonly used patterns

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

integration/control.go

@@ -1,6 +1,8 @@
 package integration
 
-import v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
+import (
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
+)
 
 type ControlServer interface {
 	Shutdown() error
@@ -9,5 +11,5 @@ type ControlServer interface {
 	WaitForReady() error
 	CreateNamespace(namespace string) error
 	CreateAuthKey(namespace string) (*v1.PreAuthKey, error)
-	ListNodes(namespace string) ([]*v1.Machine, error)
+	ListMachinesInNamespace(namespace string) ([]*v1.Machine, error)
 }

integration/general_test.go

@@ -1,7 +1,6 @@
 package integration
 
 import (
-	"net/netip"
 	"testing"
 )
 
@@ -23,30 +22,14 @@ func TestPingAllByIP(t *testing.T) {
 		t.Errorf("failed to create headscale environment: %s", err)
 	}
 
-	var allIps []netip.Addr
-	var allClients []TailscaleClient
+	allClients, err := scenario.ListTailscaleClients()
+	if err != nil {
+		t.Errorf("failed to get clients: %s", err)
+	}
 
-	for namespace, count := range spec {
-		ips, err := scenario.GetIPs(namespace)
-		if err != nil {
-			t.Errorf("failed to get tailscale ips: %s", err)
-		}
-
-		if len(ips) != count*2 {
-			t.Errorf(
-				"got the wrong amount of tailscale ips, %d != %d",
-				len(ips),
-				count*2,
-			)
-		}
-
-		clients, err := scenario.GetClients(namespace)
-		if err != nil {
-			t.Errorf("failed to get tailscale clients: %s", err)
-		}
-
-		allIps = append(allIps, ips...)
-		allClients = append(allClients, clients...)
+	allIps, err := scenario.ListTailscaleClientsIPs()
+	if err != nil {
+		t.Errorf("failed to get clients: %s", err)
 	}
 
 	err = scenario.WaitForTailscaleSync()
@@ -94,16 +77,9 @@ func TestPingAllByHostname(t *testing.T) {
 		t.Errorf("failed to create headscale environment: %s", err)
 	}
 
-	allClients := make([]TailscaleClient, 0)
-	allHostnames := make([]string, 0)
-
-	for namespace := range spec {
-		clients, err := scenario.GetClients(namespace)
-		if err != nil {
-			t.Errorf("failed to get tailscale clients: %s", err)
-		}
-
-		allClients = append(allClients, clients...)
+	allClients, err := scenario.ListTailscaleClients()
+	if err != nil {
+		t.Errorf("failed to get clients: %s", err)
 	}
 
 	err = scenario.WaitForTailscaleSync()
@@ -111,13 +87,9 @@ func TestPingAllByHostname(t *testing.T) {
 		t.Errorf("failed wait for tailscale clients to be in sync: %s", err)
 	}
 
-	for _, client := range allClients {
-		fqdn, err := client.FQDN()
-		if err != nil {
-			t.Errorf("failed to get fqdn of client %s: %s", client.Hostname(), err)
-		}
-
-		allHostnames = append(allHostnames, fqdn)
+	allHostnames, err := scenario.ListTailscaleClientsFQDNs()
+	if err != nil {
+		t.Errorf("failed to get FQDNs: %s", err)
 	}
 
 	success := 0

integration/hsic/hsic.go

@@ -199,7 +199,7 @@ func (t *HeadscaleInContainer) CreateAuthKey(
 	return &preAuthKey, nil
 }
 
-func (t *HeadscaleInContainer) ListNodes(
+func (t *HeadscaleInContainer) ListMachinesInNamespace(
 	namespace string,
 ) ([]*v1.Machine, error) {
 	command := []string{"headscale", "--namespace", namespace, "nodes", "list", "--output", "json"}

integration/scenario.go

@@ -138,6 +138,15 @@ func (s *Scenario) Shutdown() error {
 	return nil
 }
 
+func (s *Scenario) Namespaces() []string {
+	namespaces := make([]string, 0)
+	for namespace := range s.namespaces {
+		namespaces = append(namespaces, namespace)
+	}
+
+	return namespaces
+}
+
 /// Headscale related stuff
 // Note: These functions assume that there is a _single_ headscale instance for now
 
@@ -345,3 +354,61 @@ func (s *Scenario) GetClients(namespace string) ([]TailscaleClient, error) {
 
 	return clients, fmt.Errorf("failed to get clients: %w", errNoNamespaceAvailable)
 }
+
+func (s *Scenario) ListTailscaleClients(namespaces ...string) ([]TailscaleClient, error) {
+	var allClients []TailscaleClient
+
+	if len(namespaces) == 0 {
+		namespaces = s.Namespaces()
+	}
+
+	for _, namespace := range namespaces {
+		clients, err := s.GetClients(namespace)
+		if err != nil {
+			return nil, err
+		}
+
+		allClients = append(allClients, clients...)
+	}
+
+	return allClients, nil
+}
+
+func (s *Scenario) ListTailscaleClientsIPs(namespaces ...string) ([]netip.Addr, error) {
+	var allIps []netip.Addr
+
+	if len(namespaces) == 0 {
+		namespaces = s.Namespaces()
+	}
+
+	for _, namespace := range namespaces {
+		ips, err := s.GetIPs(namespace)
+		if err != nil {
+			return nil, err
+		}
+
+		allIps = append(allIps, ips...)
+	}
+
+	return allIps, nil
+}
+
+func (s *Scenario) ListTailscaleClientsFQDNs(namespaces ...string) ([]string, error) {
+	allFQDNs := make([]string, 0)
+
+	clients, err := s.ListTailscaleClients(namespaces...)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, client := range clients {
+		fqdn, err := client.FQDN()
+		if err != nil {
+			return nil, err
+		}
+
+		allFQDNs = append(allFQDNs, fqdn)
+	}
+
+	return allFQDNs, nil
+}