From f402a87e5d04fdf1c6f24fc89380f0f67d1a87c4 Mon Sep 17 00:00:00 2001
From: phm07 <22707808+phm07@users.noreply.github.com>
Date: Mon, 26 Aug 2024 16:07:57 +0200
Subject: [PATCH] fix(primary-ip): change protection commands do not allow
 protection levels

---
 internal/cmd/primaryip/disable_protection.go  |  2 ++
 .../cmd/primaryip/disable_protection_test.go  | 32 ++++++++++++++++++-
 internal/cmd/primaryip/enable_protection.go   |  2 ++
 .../cmd/primaryip/enable_protection_test.go   | 29 +++++++++++++++++
 4 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/internal/cmd/primaryip/disable_protection.go b/internal/cmd/primaryip/disable_protection.go
index ebbe4cbe..f03c8340 100644
--- a/internal/cmd/primaryip/disable_protection.go
+++ b/internal/cmd/primaryip/disable_protection.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/hetznercloud/cli/internal/cmd/base"
 	"github.com/hetznercloud/cli/internal/cmd/cmpl"
+	"github.com/hetznercloud/cli/internal/cmd/util"
 	"github.com/hetznercloud/cli/internal/hcapi2"
 	"github.com/hetznercloud/cli/internal/state"
 )
@@ -15,6 +16,7 @@ var DisableProtectionCmd = base.Cmd{
 	BaseCobraCommand: func(client hcapi2.Client) *cobra.Command {
 		cmd := &cobra.Command{
 			Use:   "disable-protection <primary-ip> [<protection-level>...]", // optional because of backwards compatibility
+			Args:  util.ValidateLenient,
 			Short: "Disable Protection for a Primary IP",
 			ValidArgsFunction: cmpl.SuggestArgs(
 				cmpl.SuggestCandidatesF(client.PrimaryIP().Names),
diff --git a/internal/cmd/primaryip/disable_protection_test.go b/internal/cmd/primaryip/disable_protection_test.go
index e6e9c927..531bfd13 100644
--- a/internal/cmd/primaryip/disable_protection_test.go
+++ b/internal/cmd/primaryip/disable_protection_test.go
@@ -11,7 +11,7 @@ import (
 	"github.com/hetznercloud/hcloud-go/v2/hcloud"
 )
 
-func TestEnable(t *testing.T) {
+func TestDisableProtection(t *testing.T) {
 	fx := testutil.NewFixture(t)
 	defer fx.Finish()
 
@@ -52,3 +52,33 @@ func TestEnable(t *testing.T) {
 	assert.Empty(t, errOut)
 	assert.Equal(t, expOut, out)
 }
+
+func TestDisableDeleteProtection(t *testing.T) {
+	fx := testutil.NewFixture(t)
+	defer fx.Finish()
+
+	cmd := primaryip.DisableProtectionCmd.CobraCommand(fx.State())
+	action := &hcloud.Action{ID: 1}
+	ip := &hcloud.PrimaryIP{ID: 13}
+
+	fx.ExpectEnsureToken()
+	fx.Client.PrimaryIPClient.EXPECT().
+		Get(gomock.Any(), "13").
+		Return(ip, nil, nil)
+	fx.Client.PrimaryIPClient.EXPECT().
+		ChangeProtection(gomock.Any(), hcloud.PrimaryIPChangeProtectionOpts{
+			ID:     13,
+			Delete: false,
+		}).
+		Return(action, nil, nil)
+
+	fx.ActionWaiter.EXPECT().WaitForActions(gomock.Any(), gomock.Any(), action).Return(nil)
+
+	out, errOut, err := fx.Run(cmd, []string{"13", "delete"})
+
+	expOut := "Resource protection disabled for primary IP 13\n"
+
+	assert.NoError(t, err)
+	assert.Empty(t, errOut)
+	assert.Equal(t, expOut, out)
+}
diff --git a/internal/cmd/primaryip/enable_protection.go b/internal/cmd/primaryip/enable_protection.go
index f8cc6038..9cf3be98 100644
--- a/internal/cmd/primaryip/enable_protection.go
+++ b/internal/cmd/primaryip/enable_protection.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/hetznercloud/cli/internal/cmd/base"
 	"github.com/hetznercloud/cli/internal/cmd/cmpl"
+	"github.com/hetznercloud/cli/internal/cmd/util"
 	"github.com/hetznercloud/cli/internal/hcapi2"
 	"github.com/hetznercloud/cli/internal/state"
 	"github.com/hetznercloud/hcloud-go/v2/hcloud"
@@ -57,6 +58,7 @@ var EnableProtectionCmd = base.Cmd{
 	BaseCobraCommand: func(client hcapi2.Client) *cobra.Command {
 		cmd := &cobra.Command{
 			Use:   "enable-protection <primary-ip> [<protection-level>...]", // optional because of backwards compatibility
+			Args:  util.ValidateLenient,
 			Short: "Enable Protection for a Primary IP",
 			ValidArgsFunction: cmpl.SuggestArgs(
 				cmpl.SuggestCandidatesF(client.PrimaryIP().Names),
diff --git a/internal/cmd/primaryip/enable_protection_test.go b/internal/cmd/primaryip/enable_protection_test.go
index a780b757..559d98f1 100644
--- a/internal/cmd/primaryip/enable_protection_test.go
+++ b/internal/cmd/primaryip/enable_protection_test.go
@@ -52,3 +52,32 @@ func TestEnableProtection(t *testing.T) {
 	assert.Empty(t, errOut)
 	assert.Equal(t, expOut, out)
 }
+
+func TestEnableDeleteProtection(t *testing.T) {
+	fx := testutil.NewFixture(t)
+	defer fx.Finish()
+
+	cmd := primaryip.EnableProtectionCmd.CobraCommand(fx.State())
+	action := &hcloud.Action{ID: 1}
+	ip := &hcloud.PrimaryIP{ID: 13}
+
+	fx.ExpectEnsureToken()
+	fx.Client.PrimaryIPClient.EXPECT().
+		Get(gomock.Any(), "13").
+		Return(ip, nil, nil)
+	fx.Client.PrimaryIPClient.EXPECT().
+		ChangeProtection(gomock.Any(), hcloud.PrimaryIPChangeProtectionOpts{
+			ID:     13,
+			Delete: true,
+		}).
+		Return(action, nil, nil)
+	fx.ActionWaiter.EXPECT().WaitForActions(gomock.Any(), gomock.Any(), action).Return(nil)
+
+	out, errOut, err := fx.Run(cmd, []string{"13", "delete"})
+
+	expOut := "Resource protection enabled for primary IP 13\n"
+
+	assert.NoError(t, err)
+	assert.Empty(t, errOut)
+	assert.Equal(t, expOut, out)
+}