From fd5373785faab3b72ec8867c03902cad2cb9b9b8 Mon Sep 17 00:00:00 2001
From: Suvorov Daniil <severecloud@gmail.com>
Date: Wed, 19 May 2021 20:07:05 +0300
Subject: [PATCH] feat: add callback remove funcion

---
 callback/callback.go     | 16 ++++++++++++++--
 callback/context.go      |  5 +++++
 callback/context_test.go | 23 +++++++++++++++++++++++
 internal/transport.go    |  1 +
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/callback/callback.go b/callback/callback.go
index ff7332ba..08e1277c 100644
--- a/callback/callback.go
+++ b/callback/callback.go
@@ -83,8 +83,9 @@ func (cb *Callback) HandleFunc(w http.ResponseWriter, r *http.Request) {
 	ctx = context.WithValue(ctx, internal.CallbackRetryCounterKey, retryCounter)
 
 	var (
-		code int
-		date time.Time
+		code   int
+		date   time.Time
+		remove bool
 	)
 
 	retryAfter := func(c int, d time.Time) {
@@ -93,6 +94,11 @@ func (cb *Callback) HandleFunc(w http.ResponseWriter, r *http.Request) {
 	}
 	ctx = context.WithValue(ctx, internal.CallbackRetryAfterKey, retryAfter)
 
+	removeFunc := func() {
+		remove = true
+	}
+	ctx = context.WithValue(ctx, internal.CallbackRemove, removeFunc)
+
 	if err := cb.Handler(ctx, e); err != nil {
 		cb.logf("callback: %v", err)
 		http.Error(w, "Bad Request", http.StatusBadRequest)
@@ -100,6 +106,12 @@ func (cb *Callback) HandleFunc(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if remove {
+		_, _ = w.Write([]byte("remove"))
+
+		return
+	}
+
 	if code != 0 {
 		w.Header().Set("Retry-After", date.Format(http.TimeFormat)) // RFC 7231, 7.1.3
 		http.Error(w, http.StatusText(code), code)
diff --git a/callback/context.go b/callback/context.go
index 9412a8f4..10f67400 100644
--- a/callback/context.go
+++ b/callback/context.go
@@ -30,3 +30,8 @@ func RetryAfter(ctx context.Context, code int, date time.Time) {
 		date,
 	)
 }
+
+// Remove VK Callback server.
+func Remove(ctx context.Context) {
+	ctx.Value(internal.CallbackRemove).(func())()
+}
diff --git a/callback/context_test.go b/callback/context_test.go
index d4fe0dd9..b9ec6fa3 100644
--- a/callback/context_test.go
+++ b/callback/context_test.go
@@ -57,3 +57,26 @@ func TestRetryAfter(t *testing.T) {
 	assert.Equal(t, code, rr.Code)
 	assert.Equal(t, date.Format(http.TimeFormat), rr.Header().Get("Retry-After"))
 }
+
+func TestRemove(t *testing.T) {
+	t.Parallel()
+
+	cb := callback.NewCallback()
+	cb.MessageNew(func(ctx context.Context, obj events.MessageNewObject) {
+		callback.Remove(ctx)
+	})
+
+	jsonStr := []byte(`{"type": "message_new","object": {}}`)
+
+	req, err := http.NewRequest("POST", "/callback", bytes.NewBuffer(jsonStr))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rr := httptest.NewRecorder()
+	handler := http.HandlerFunc(cb.HandleFunc)
+
+	handler.ServeHTTP(rr, req)
+
+	assert.Equal(t, "remove", rr.Body.String())
+}
diff --git a/internal/transport.go b/internal/transport.go
index 2ff8fbb3..6611b496 100644
--- a/internal/transport.go
+++ b/internal/transport.go
@@ -27,6 +27,7 @@ const (
 	LongPollTsKey
 	CallbackRetryCounterKey
 	CallbackRetryAfterKey
+	CallbackRemove
 )
 
 // ContextClient return *http.Client.