Define a "scenario", which is a controlserver with nodes

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

integration/scenario.go

@@ -0,0 +1,206 @@
+package integration
+
+import (
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"sync"
+
+	"github.com/juanfont/headscale"
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
+	"github.com/juanfont/headscale/integration/dockertestutil"
+	"github.com/juanfont/headscale/integration/hsic"
+	"github.com/juanfont/headscale/integration/tsic"
+	"github.com/ory/dockertest/v3"
+)
+
+const scenarioHashLength = 6
+
+var errNoHeadscaleAvailable = errors.New("no headscale available")
+var errNoNamespaceAvailable = errors.New("no namespace available")
+
+type Namespace struct {
+	Clients map[string]*tsic.TailscaleInContainer
+
+	createWaitGroup sync.WaitGroup
+	joinWaitGroup   sync.WaitGroup
+}
+
+// TODO(kradalby): make control server configurable, test test correctness with
+// Tailscale SaaS.
+type Scenario struct {
+	// TODO(kradalby): support multiple headcales for later, currently only
+	// use one.
+	controlServers map[string]ControlServer
+
+	namespaces map[string]*Namespace
+
+	pool    *dockertest.Pool
+	network *dockertest.Network
+}
+
+func NewScenario() (*Scenario, error) {
+	hash, err := headscale.GenerateRandomStringDNSSafe(scenarioHashLength)
+	if err != nil {
+		return nil, err
+	}
+
+	pool, err := dockertest.NewPool("")
+	if err != nil {
+		return nil, fmt.Errorf("could not connect to docker: %w", err)
+	}
+
+	networkName := fmt.Sprintf("hs-%s", hash)
+	if overrideNetworkName := os.Getenv("HEADSCALE_TEST_NETWORK_NAME"); overrideNetworkName != "" {
+		networkName = overrideNetworkName
+	}
+
+	network, err := dockertestutil.GetFirstOrCreateNetwork(pool, networkName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create or get network: %w", err)
+	}
+
+	return &Scenario{
+		controlServers: make(map[string]ControlServer),
+		namespaces:     make(map[string]*Namespace),
+
+		pool:    pool,
+		network: network,
+	}, nil
+}
+
+func (s *Scenario) Shutdown() error {
+	for _, control := range s.controlServers {
+		err := control.Shutdown()
+		if err != nil {
+			return fmt.Errorf("failed to tear down control: %w", err)
+		}
+	}
+
+	for namespaceName, namespace := range s.namespaces {
+		for _, client := range namespace.Clients {
+			log.Printf("removing client %s in namespace %s", client.Hostname, namespaceName)
+			err := client.Shutdown()
+			if err != nil {
+				return fmt.Errorf("failed to tear down client: %w", err)
+			}
+		}
+	}
+
+	// TODO(kradalby): This breaks the "we need to create a network before we start"
+	// part, since we now run the tests in a container...
+	// if err := s.pool.RemoveNetwork(s.network); err != nil {
+	// 	return fmt.Errorf("failed to remove network: %w", err)
+	// }
+
+	// TODO(kradalby): This seem redundant to the previous call
+	// if err := s.network.Close(); err != nil {
+	// 	return fmt.Errorf("failed to tear down network: %w", err)
+	// }
+
+	return nil
+}
+
+/// Headscale related stuff
+// Note: These functions assume that there is a _single_ headscale instance for now
+
+// TODO(kradalby): make port and headscale configurable, multiple instances support?
+func (s *Scenario) StartHeadscale() error {
+	headscale, err := hsic.New(s.pool, 8080, s.network)
+	if err != nil {
+		return fmt.Errorf("failed to create headscale container: %w", err)
+	}
+
+	s.controlServers["headscale"] = headscale
+
+	return nil
+}
+
+func (s *Scenario) Headscale() *hsic.HeadscaleInContainer {
+	return s.controlServers["headscale"].(*hsic.HeadscaleInContainer)
+}
+
+func (s *Scenario) CreatePreAuthKey(namespace string) (*v1.PreAuthKey, error) {
+	if headscale, ok := s.controlServers["headscale"]; ok {
+		key, err := headscale.CreateAuthKey(namespace)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create namespace: %w", err)
+		}
+
+		return key, nil
+	}
+
+	return nil, fmt.Errorf("failed to create namespace: %w", errNoHeadscaleAvailable)
+}
+
+func (s *Scenario) CreateNamespace(namespace string) error {
+	if headscale, ok := s.controlServers["headscale"]; ok {
+		err := headscale.CreateNamespace(namespace)
+		if err != nil {
+			return fmt.Errorf("failed to create namespace: %w", err)
+		}
+
+		s.namespaces[namespace] = &Namespace{
+			Clients: make(map[string]*tsic.TailscaleInContainer),
+		}
+
+		return nil
+	}
+
+	return fmt.Errorf("failed to create namespace: %w", errNoHeadscaleAvailable)
+}
+
+/// Client related stuff
+
+func (s *Scenario) CreateTailscaleNodesInNamespace(
+	namespace string,
+	version string,
+	count int,
+) error {
+	if ns, ok := s.namespaces[namespace]; ok {
+		for i := 0; i < count; i++ {
+			ns.createWaitGroup.Add(1)
+
+			go func() {
+				defer ns.createWaitGroup.Done()
+
+				// TODO(kradalby): error handle this
+				ts, err := tsic.New(s.pool, version, s.network)
+				if err != nil {
+					// return fmt.Errorf("failed to add tailscale node: %w", err)
+					fmt.Printf("failed to add tailscale node: %s", err)
+				}
+
+				ns.Clients[ts.Hostname] = ts
+			}()
+		}
+		ns.createWaitGroup.Wait()
+
+		return nil
+	}
+
+	return fmt.Errorf("failed to add tailscale node: %w", errNoNamespaceAvailable)
+}
+
+func (s *Scenario) RunTailscaleUp(
+	namespace, loginServer, authKey string,
+) error {
+	if ns, ok := s.namespaces[namespace]; ok {
+		for _, client := range ns.Clients {
+			ns.joinWaitGroup.Add(1)
+
+			go func() {
+				defer ns.joinWaitGroup.Done()
+
+				// TODO(kradalby): error handle this
+				_ = client.Up(loginServer, authKey)
+			}()
+		}
+		ns.joinWaitGroup.Wait()
+
+		return nil
+	}
+
+	return fmt.Errorf("failed to up tailscale node: %w", errNoNamespaceAvailable)
+}

integration/scenario_test.go

@@ -0,0 +1,149 @@
+package integration
+
+import (
+	"testing"
+
+	"github.com/juanfont/headscale/integration/tsic"
+)
+
+func TestHeadscale(t *testing.T) {
+	var err error
+
+	scenario, err := NewScenario()
+	if err != nil {
+		t.Errorf("failed to create scenario: %s", err)
+	}
+
+	t.Run("start-headscale", func(t *testing.T) {
+		err = scenario.StartHeadscale()
+		if err != nil {
+			t.Errorf("failed to create start headcale: %s", err)
+		}
+
+		err = scenario.Headscale().WaitForReady()
+		if err != nil {
+			t.Errorf("headscale failed to become ready: %s", err)
+		}
+
+	})
+
+	t.Run("create-namespace", func(t *testing.T) {
+		err := scenario.CreateNamespace("test-space")
+		if err != nil {
+			t.Errorf("failed to create namespace: %s", err)
+		}
+
+		if _, ok := scenario.namespaces["test-space"]; !ok {
+			t.Errorf("namespace is not in scenario")
+		}
+	})
+
+	t.Run("create-auth-key", func(t *testing.T) {
+		_, err := scenario.CreatePreAuthKey("test-space")
+		if err != nil {
+			t.Errorf("failed to create preauthkey: %s", err)
+		}
+	})
+
+	err = scenario.Shutdown()
+	if err != nil {
+		t.Errorf("failed to tear down scenario: %s", err)
+	}
+}
+
+func TestCreateTailscale(t *testing.T) {
+	var scenario *Scenario
+	var err error
+
+	namespace := "only-create-containers"
+
+	scenario, err = NewScenario()
+	if err != nil {
+		t.Errorf("failed to create scenario: %s", err)
+	}
+
+	scenario.namespaces[namespace] = &Namespace{
+		Clients: make(map[string]*tsic.TailscaleInContainer),
+	}
+
+	t.Run("create-tailscale", func(t *testing.T) {
+		err := scenario.CreateTailscaleNodesInNamespace(namespace, "1.32.0", 3)
+		if err != nil {
+			t.Errorf("failed to add tailscale nodes: %s", err)
+		}
+
+		if clients := len(scenario.namespaces[namespace].Clients); clients != 3 {
+			t.Errorf("wrong number of tailscale clients: %d != %d", clients, 3)
+		}
+	})
+
+	err = scenario.Shutdown()
+	if err != nil {
+		t.Errorf("failed to tear down scenario: %s", err)
+	}
+}
+
+func TestTailscaleNodesJoiningHeadcale(t *testing.T) {
+	var err error
+
+	namespace := "join-node-test"
+
+	scenario, err := NewScenario()
+	if err != nil {
+		t.Errorf("failed to create scenario: %s", err)
+	}
+
+	t.Run("start-headscale", func(t *testing.T) {
+		err = scenario.StartHeadscale()
+		if err != nil {
+			t.Errorf("failed to create start headcale: %s", err)
+		}
+
+		headscale := scenario.Headscale()
+		err = headscale.WaitForReady()
+		if err != nil {
+			t.Errorf("headscale failed to become ready: %s", err)
+		}
+
+	})
+
+	t.Run("create-namespace", func(t *testing.T) {
+		err := scenario.CreateNamespace(namespace)
+		if err != nil {
+			t.Errorf("failed to create namespace: %s", err)
+		}
+
+		if _, ok := scenario.namespaces[namespace]; !ok {
+			t.Errorf("namespace is not in scenario")
+		}
+	})
+
+	t.Run("create-tailscale", func(t *testing.T) {
+		err := scenario.CreateTailscaleNodesInNamespace(namespace, "1.32.0", 2)
+		if err != nil {
+			t.Errorf("failed to add tailscale nodes: %s", err)
+		}
+
+		if clients := len(scenario.namespaces[namespace].Clients); clients != 2 {
+			t.Errorf("wrong number of tailscale clients: %d != %d", clients, 2)
+		}
+	})
+
+	t.Run("join-headscale", func(t *testing.T) {
+		key, err := scenario.CreatePreAuthKey(namespace)
+		if err != nil {
+			t.Errorf("failed to create preauthkey: %s", err)
+		}
+
+		err = scenario.RunTailscaleUp(namespace, scenario.Headscale().GetEndpoint(), key.GetKey())
+		if err != nil {
+			t.Errorf("failed to login: %s", err)
+		}
+
+	})
+
+	err = scenario.Shutdown()
+	if err != nil {
+		t.Errorf("failed to tear down scenario: %s", err)
+	}
+}