From 900aced7d5e275b807617fcf915435025a85c441 Mon Sep 17 00:00:00 2001
From: Jian Yuan Lee <jianyuan@gmail.com>
Date: Thu, 7 Dec 2023 01:31:45 +0000
Subject: [PATCH] feat: Enable/disable spike protection

---
 sentry/sentry.go                 | 12 ++++----
 sentry/spike_protections.go      | 33 ++++++++++++++++++++++
 sentry/spike_protections_test.go | 47 ++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 5 deletions(-)
 create mode 100644 sentry/spike_protections.go
 create mode 100644 sentry/spike_protections_test.go

diff --git a/sentry/sentry.go b/sentry/sentry.go
index 19bfd63..ca249fb 100644
--- a/sentry/sentry.go
+++ b/sentry/sentry.go
@@ -47,8 +47,8 @@ type Client struct {
 	common service
 
 	// Services
-	DashboardWidgets         *DashboardWidgetsService
 	Dashboards               *DashboardsService
+	DashboardWidgets         *DashboardWidgetsService
 	IssueAlerts              *IssueAlertsService
 	MetricAlerts             *MetricAlertsService
 	OrganizationCodeMappings *OrganizationCodeMappingsService
@@ -56,14 +56,15 @@ type Client struct {
 	OrganizationMembers      *OrganizationMembersService
 	OrganizationRepositories *OrganizationRepositoriesService
 	Organizations            *OrganizationsService
+	ProjectFilters           *ProjectFiltersService
 	ProjectKeys              *ProjectKeysService
 	ProjectOwnerships        *ProjectOwnershipsService
 	ProjectPlugins           *ProjectPluginsService
 	Projects                 *ProjectsService
-	ProjectFilters           *ProjectFiltersService
 	ReleaseDeployments       *ReleaseDeploymentsService
-	Teams                    *TeamsService
+	SpikeProtections         *SpikeProtectionsService
 	TeamMembers              *TeamMembersService
+	Teams                    *TeamsService
 }
 
 type service struct {
@@ -84,8 +85,8 @@ func NewClient(httpClient *http.Client) *Client {
 		UserAgent: userAgent,
 	}
 	c.common.client = c
-	c.DashboardWidgets = (*DashboardWidgetsService)(&c.common)
 	c.Dashboards = (*DashboardsService)(&c.common)
+	c.DashboardWidgets = (*DashboardWidgetsService)(&c.common)
 	c.IssueAlerts = (*IssueAlertsService)(&c.common)
 	c.MetricAlerts = (*MetricAlertsService)(&c.common)
 	c.OrganizationCodeMappings = (*OrganizationCodeMappingsService)(&c.common)
@@ -99,8 +100,9 @@ func NewClient(httpClient *http.Client) *Client {
 	c.ProjectPlugins = (*ProjectPluginsService)(&c.common)
 	c.Projects = (*ProjectsService)(&c.common)
 	c.ReleaseDeployments = (*ReleaseDeploymentsService)(&c.common)
-	c.Teams = (*TeamsService)(&c.common)
+	c.SpikeProtections = (*SpikeProtectionsService)(&c.common)
 	c.TeamMembers = (*TeamMembersService)(&c.common)
+	c.Teams = (*TeamsService)(&c.common)
 	return c
 }
 
diff --git a/sentry/spike_protections.go b/sentry/spike_protections.go
new file mode 100644
index 0000000..20c0d6b
--- /dev/null
+++ b/sentry/spike_protections.go
@@ -0,0 +1,33 @@
+package sentry
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+)
+
+type SpikeProtectionsService service
+
+type SpikeProtectionParams struct {
+	Projects []string `json:"projects"`
+}
+
+func (s *SpikeProtectionsService) Enable(ctx context.Context, organizationSlug string, params *SpikeProtectionParams) (*Response, error) {
+	u := fmt.Sprintf("0/organizations/%v/spike-protections/", organizationSlug)
+	req, err := s.client.NewRequest(http.MethodPost, u, params)
+	if err != nil {
+		return nil, err
+	}
+
+	return s.client.Do(ctx, req, nil)
+}
+
+func (s *SpikeProtectionsService) Disable(ctx context.Context, organizationSlug string, params *SpikeProtectionParams) (*Response, error) {
+	u := fmt.Sprintf("0/organizations/%v/spike-protections/", organizationSlug)
+	req, err := s.client.NewRequest(http.MethodDelete, u, params)
+	if err != nil {
+		return nil, err
+	}
+
+	return s.client.Do(ctx, req, nil)
+}
diff --git a/sentry/spike_protections_test.go b/sentry/spike_protections_test.go
new file mode 100644
index 0000000..ba9b012
--- /dev/null
+++ b/sentry/spike_protections_test.go
@@ -0,0 +1,47 @@
+package sentry
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSpikeProtectionsService_Enable(t *testing.T) {
+	client, mux, _, teardown := setup()
+	defer teardown()
+
+	mux.HandleFunc("/api/0/organizations/organization_slug/spike-protections/", func(w http.ResponseWriter, r *http.Request) {
+		assertMethod(t, http.MethodPost, r)
+		assertPostJSON(t, map[string]interface{}{
+			"projects": []interface{}{"$all"},
+		}, r)
+	})
+
+	params := &SpikeProtectionParams{
+		Projects: []string{"$all"},
+	}
+	ctx := context.Background()
+	_, err := client.SpikeProtections.Enable(ctx, "organization_slug", params)
+	assert.NoError(t, err)
+}
+
+func TestSpikeProtectionsService_Disable(t *testing.T) {
+	client, mux, _, teardown := setup()
+	defer teardown()
+
+	mux.HandleFunc("/api/0/organizations/organization_slug/spike-protections/", func(w http.ResponseWriter, r *http.Request) {
+		assertMethod(t, http.MethodDelete, r)
+		assertPostJSON(t, map[string]interface{}{
+			"projects": []interface{}{"$all"},
+		}, r)
+	})
+
+	params := &SpikeProtectionParams{
+		Projects: []string{"$all"},
+	}
+	ctx := context.Background()
+	_, err := client.SpikeProtections.Disable(ctx, "organization_slug", params)
+	assert.NoError(t, err)
+}