From 239ef16ad1fd97611ceac8e76172ac988411f35a Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Mon, 24 Oct 2022 16:42:09 +0200
Subject: [PATCH] Add preauthkey command test

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 integration/cli_test.go | 139 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 139 insertions(+)

diff --git a/integration/cli_test.go b/integration/cli_test.go
index c934686..b0d9906 100644
--- a/integration/cli_test.go
+++ b/integration/cli_test.go
@@ -3,6 +3,7 @@ package integration
 import (
 	"encoding/json"
 	"testing"
+	"time"
 
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/stretchr/testify/assert"
@@ -91,3 +92,141 @@ func TestNamespaceCommand(t *testing.T) {
 	err = scenario.Shutdown()
 	assert.NoError(t, err)
 }
+
+func TestPreAuthKeyCommand(t *testing.T) {
+	IntegrationSkip(t)
+	t.Parallel()
+
+	namespace := "preauthkeyspace"
+	count := 3
+
+	scenario, err := NewScenario()
+	assert.NoError(t, err)
+
+	spec := map[string]int{
+		namespace: 0,
+	}
+
+	err = scenario.CreateHeadscaleEnv(spec)
+	assert.NoError(t, err)
+
+	keys := make([]*v1.PreAuthKey, count)
+	assert.NoError(t, err)
+
+	for index := 0; index < count; index++ {
+		var preAuthKey v1.PreAuthKey
+		err := executeAndUnmarshal(
+			scenario.Headscale(),
+			[]string{
+				"headscale",
+				"preauthkeys",
+				"--namespace",
+				namespace,
+				"create",
+				"--reusable",
+				"--expiration",
+				"24h",
+				"--output",
+				"json",
+				"--tags",
+				"tag:test1,tag:test2",
+			},
+			&preAuthKey,
+		)
+		assert.NoError(t, err)
+
+		keys[index] = &preAuthKey
+	}
+
+	assert.Len(t, keys, 3)
+
+	var listedPreAuthKeys []v1.PreAuthKey
+	err = executeAndUnmarshal(
+		scenario.Headscale(),
+		[]string{
+			"headscale",
+			"preauthkeys",
+			"--namespace",
+			namespace,
+			"list",
+			"--output",
+			"json",
+		},
+		&listedPreAuthKeys,
+	)
+	assert.NoError(t, err)
+
+	// There is one key created by "scenario.CreateHeadscaleEnv"
+	assert.Len(t, listedPreAuthKeys, 4)
+
+	assert.Equal(
+		t,
+		[]string{keys[0].Id, keys[1].Id, keys[2].Id},
+		[]string{listedPreAuthKeys[1].Id, listedPreAuthKeys[2].Id, listedPreAuthKeys[3].Id},
+	)
+
+	assert.NotEmpty(t, listedPreAuthKeys[1].Key)
+	assert.NotEmpty(t, listedPreAuthKeys[2].Key)
+	assert.NotEmpty(t, listedPreAuthKeys[3].Key)
+
+	assert.True(t, listedPreAuthKeys[1].Expiration.AsTime().After(time.Now()))
+	assert.True(t, listedPreAuthKeys[2].Expiration.AsTime().After(time.Now()))
+	assert.True(t, listedPreAuthKeys[3].Expiration.AsTime().After(time.Now()))
+
+	assert.True(
+		t,
+		listedPreAuthKeys[1].Expiration.AsTime().Before(time.Now().Add(time.Hour*26)),
+	)
+	assert.True(
+		t,
+		listedPreAuthKeys[2].Expiration.AsTime().Before(time.Now().Add(time.Hour*26)),
+	)
+	assert.True(
+		t,
+		listedPreAuthKeys[3].Expiration.AsTime().Before(time.Now().Add(time.Hour*26)),
+	)
+
+	for index := range listedPreAuthKeys {
+		if index == 0 {
+			continue
+		}
+
+		assert.Equal(t, listedPreAuthKeys[index].AclTags, []string{"tag:test1", "tag:test2"})
+	}
+
+	// Test key expiry
+	_, err = scenario.Headscale().Execute(
+		[]string{
+			"headscale",
+			"preauthkeys",
+			"--namespace",
+			namespace,
+			"expire",
+			listedPreAuthKeys[1].Key,
+		},
+	)
+	assert.NoError(t, err)
+
+	var listedPreAuthKeysAfterExpire []v1.PreAuthKey
+	err = executeAndUnmarshal(
+		scenario.Headscale(),
+		[]string{
+			"headscale",
+			"preauthkeys",
+			"--namespace",
+			namespace,
+			"list",
+			"--output",
+			"json",
+		},
+		&listedPreAuthKeysAfterExpire,
+	)
+	assert.NoError(t, err)
+
+	assert.True(t, listedPreAuthKeysAfterExpire[1].Expiration.AsTime().Before(time.Now()))
+	assert.True(t, listedPreAuthKeysAfterExpire[2].Expiration.AsTime().After(time.Now()))
+	assert.True(t, listedPreAuthKeysAfterExpire[3].Expiration.AsTime().After(time.Now()))
+
+	err = scenario.Shutdown()
+	assert.NoError(t, err)
+}