From 5a0be76053b9d8a485d9ca100e02fbc7e33c1f38 Mon Sep 17 00:00:00 2001
From: Pallab Pain <pallabkumarpain@gmail.com>
Date: Sat, 13 Jul 2024 00:30:33 +0530
Subject: [PATCH] test: adds integration test for policy command

---
 .github/workflows/test-integration.yaml |  1 +
 integration/cli_test.go                 | 80 +++++++++++++++++++++++++
 integration/control.go                  |  4 +-
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test-integration.yaml b/.github/workflows/test-integration.yaml
index 9581badaef9..cdceb5b4e48 100644
--- a/.github/workflows/test-integration.yaml
+++ b/.github/workflows/test-integration.yaml
@@ -35,6 +35,7 @@ jobs:
           - TestNodeExpireCommand
           - TestNodeRenameCommand
           - TestNodeMoveCommand
+          - TestPolicyCommand
           - TestDERPServerScenario
           - TestPingAllByIP
           - TestPingAllByIPPublicDERP
diff --git a/integration/cli_test.go b/integration/cli_test.go
index d78c7b1650c..9bc67a8937c 100644
--- a/integration/cli_test.go
+++ b/integration/cli_test.go
@@ -1597,3 +1597,83 @@ func TestNodeMoveCommand(t *testing.T) {
 
 	assert.Equal(t, node.GetUser().GetName(), "old-user")
 }
+
+func TestPolicyCommand(t *testing.T) {
+	IntegrationSkip(t)
+	t.Parallel()
+
+	scenario, err := NewScenario(dockertestMaxWait())
+	assertNoErr(t, err)
+	defer scenario.Shutdown()
+
+	spec := map[string]int{
+		"policy-user": 0,
+	}
+
+	err = scenario.CreateHeadscaleEnv(
+		spec,
+		[]tsic.Option{},
+		hsic.WithTestName("clins"),
+		hsic.WithConfigEnv(map[string]string{
+			"HEADSCALE_POLICY_MODE": "database",
+		}),
+	)
+	assertNoErr(t, err)
+
+	headscale, err := scenario.Headscale()
+	assertNoErr(t, err)
+
+	p := policy.ACLPolicy{
+		ACLs: []policy.ACL{
+			{
+				Action:       "accept",
+				Sources:      []string{"*"},
+				Destinations: []string{"*:*"},
+			},
+		},
+		TagOwners: map[string][]string{
+			"tag:exists": {"policy-user"},
+		},
+	}
+
+	pBytes, _ := json.Marshal(p)
+
+	policyFilePath := "/etc/headscale/policy.json"
+
+	err = headscale.WriteFile(policyFilePath, pBytes)
+	assertNoErr(t, err)
+
+	// No policy is present at this time.
+	// Add a new policy from a file.
+	_, err = headscale.Execute(
+		[]string{
+			"headscale",
+			"policy",
+			"set",
+			"-f",
+			policyFilePath,
+		},
+	)
+
+	assertNoErr(t, err)
+
+	// Get the current policy and check
+	// if it is the same as the one we set.
+	var output *policy.ACLPolicy
+	err = executeAndUnmarshal(
+		headscale,
+		[]string{
+			"headscale",
+			"policy",
+			"get",
+			"--output",
+			"json",
+		},
+		&output,
+	)
+	assertNoErr(t, err)
+
+	assert.Len(t, output.TagOwners, 1)
+	assert.Len(t, output.ACLs, 1)
+	assert.Equal(t, output.TagOwners["tag:exists"], []string{"policy-user"})
+}
diff --git a/integration/control.go b/integration/control.go
index f5557495a09..4260ac4b22e 100644
--- a/integration/control.go
+++ b/integration/control.go
@@ -1,8 +1,9 @@
 package integration
 
 import (
-	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/ory/dockertest/v3"
+
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 )
 
 type ControlServer interface {
@@ -10,6 +11,7 @@ type ControlServer interface {
 	SaveLog(string) error
 	SaveProfile(string) error
 	Execute(command []string) (string, error)
+	WriteFile(path string, content []byte) error
 	ConnectToNetwork(network *dockertest.Network) error
 	GetHealthEndpoint() string
 	GetEndpoint() string