From d63df8ce0e27944fe2cc37c299702196568021b9 Mon Sep 17 00:00:00 2001
From: Demian <d@aletheia.icu>
Date: Sun, 7 Jun 2020 21:41:13 +0300
Subject: [PATCH] bot: extend Edit function

---
 bot.go      | 16 ++++++++-----
 bot_test.go | 65 +++++++++++++++++++++++++++++------------------------
 2 files changed, 47 insertions(+), 34 deletions(-)

diff --git a/bot.go b/bot.go
index 38247b1..83eaee4 100644
--- a/bot.go
+++ b/bot.go
@@ -664,9 +664,11 @@ func (b *Bot) Forward(to Recipient, msg Editable, options ...interface{}) (*Mess
 //
 // Use cases:
 //
-//     b.Edit(msg, msg.Text, newMarkup)
-//     b.Edit(msg, "new <b>text</b>", tb.ModeHTML)
-//     b.Edit(msg, tb.Location{42.1337, 69.4242})
+//     b.Edit(m, m.Text, newMarkup)
+//     b.Edit(m, "new <b>text</b>", tb.ModeHTML)
+//     b.Edit(m, &tb.ReplyMarkup{...})
+//     b.Edit(m, &tb.Photo{File: ...})
+//     b.Edit(m, tb.Location{42.1337, 69.4242})
 //
 // This function will panic upon nil Editable.
 func (b *Bot) Edit(msg Editable, what interface{}, options ...interface{}) (*Message, error) {
@@ -676,6 +678,10 @@ func (b *Bot) Edit(msg Editable, what interface{}, options ...interface{}) (*Mes
 	)
 
 	switch v := what.(type) {
+	case *ReplyMarkup:
+		return b.EditReplyMarkup(msg, v)
+	case InputMedia:
+		return b.EditMedia(msg, v, options...)
 	case string:
 		method = "editMessageText"
 		params["text"] = v
@@ -782,8 +788,8 @@ func (b *Bot) EditCaption(msg Editable, caption string, options ...interface{})
 //
 // Use cases:
 //
-//     bot.EditMedia(msg, &tb.Photo{File: tb.FromDisk("chicken.jpg")})
-//     bot.EditMedia(msg, &tb.Video{File: tb.FromURL("http://video.mp4")})
+//     b.EditMedia(m, &tb.Photo{File: tb.FromDisk("chicken.jpg")})
+//     b.EditMedia(m, &tb.Video{File: tb.FromURL("http://video.mp4")})
 //
 // This function will panic upon nil Editable.
 func (b *Bot) EditMedia(msg Editable, media InputMedia, options ...interface{}) (*Message, error) {
diff --git a/bot_test.go b/bot_test.go
index 7b21fae..2a28d94 100644
--- a/bot_test.go
+++ b/bot_test.go
@@ -309,23 +309,30 @@ func TestBot(t *testing.T) {
 	_, err = b.Forward(nil, nil)
 	assert.Equal(t, ErrBadRecipient, err)
 
-	t.Run("Send(what=Sendable)", func(t *testing.T) {
-		photo := &Photo{
-			File:    File{FileID: photoID},
-			Caption: t.Name(),
-		}
+	photo := &Photo{
+		File:    File{FileID: photoID},
+		Caption: t.Name(),
+	}
+	var msg *Message
 
-		msg, err := b.Send(to, photo)
+	t.Run("Send(what=Sendable)", func(t *testing.T) {
+		msg, err = b.Send(to, photo)
 		assert.NoError(t, err)
 		assert.NotNil(t, msg.Photo)
 		assert.Equal(t, photo.Caption, msg.Caption)
+	})
 
-		msg, err = b.EditCaption(msg, "new caption")
+	t.Run("EditCaption()", func(t *testing.T) {
+		edited, err := b.EditCaption(msg, "new caption")
 		assert.NoError(t, err)
-		assert.Equal(t, "new caption", msg.Caption)
+		assert.Equal(t, "new caption", edited.Caption)
 	})
 
-	var msg *Message
+	t.Run("Edit(what=InputMedia)", func(t *testing.T) {
+		edited, err := b.Edit(msg, photo)
+		assert.NoError(t, err)
+		assert.Equal(t, msg.Photo, edited.Photo)
+	})
 
 	t.Run("Send(what=string)", func(t *testing.T) {
 		msg, err = b.Send(to, t.Name())
@@ -358,20 +365,8 @@ func TestBot(t *testing.T) {
 		assert.Error(t, err) // message is not modified
 	})
 
-	t.Run("Edit(what=Location)", func(t *testing.T) {
-		loc := &Location{Lat: 42, Lng: 69, LivePeriod: 60}
-		msg, err := b.Send(to, loc)
-		assert.NoError(t, err)
-		assert.NotNil(t, msg.Location)
-
-		loc = &Location{Lat: loc.Lng, Lng: loc.Lat}
-		msg, err = b.Edit(msg, *loc)
-		assert.NoError(t, err)
-		assert.NotNil(t, msg.Location)
-	})
-
-	t.Run("EditReplyMarkup()", func(t *testing.T) {
-		markup := &ReplyMarkup{
+	t.Run("Edit(what=ReplyMarkup)", func(t *testing.T) {
+		good := &ReplyMarkup{
 			InlineKeyboard: [][]InlineButton{
 				{{
 					Data: "btn",
@@ -379,7 +374,7 @@ func TestBot(t *testing.T) {
 				}},
 			},
 		}
-		badMarkup := &ReplyMarkup{
+		bad := &ReplyMarkup{
 			InlineKeyboard: [][]InlineButton{
 				{{
 					Data: strings.Repeat("*", 65),
@@ -388,18 +383,30 @@ func TestBot(t *testing.T) {
 			},
 		}
 
-		msg, err := b.EditReplyMarkup(msg, markup)
+		edited, err := b.Edit(msg, good)
 		assert.NoError(t, err)
-		assert.Equal(t, msg.ReplyMarkup.InlineKeyboard, markup.InlineKeyboard)
+		assert.Equal(t, edited.ReplyMarkup.InlineKeyboard, good.InlineKeyboard)
 
-		msg, err = b.EditReplyMarkup(msg, nil)
+		edited, err = b.EditReplyMarkup(edited, nil)
 		assert.NoError(t, err)
-		assert.Nil(t, msg.ReplyMarkup.InlineKeyboard)
+		assert.Nil(t, edited.ReplyMarkup.InlineKeyboard)
 
-		_, err = b.EditReplyMarkup(msg, badMarkup)
+		_, err = b.Edit(edited, bad)
 		assert.Equal(t, ErrButtonDataInvalid, err)
 	})
 
+	t.Run("Edit(what=Location)", func(t *testing.T) {
+		loc := &Location{Lat: 42, Lng: 69, LivePeriod: 60}
+		edited, err := b.Send(to, loc)
+		assert.NoError(t, err)
+		assert.NotNil(t, edited.Location)
+
+		loc = &Location{Lat: loc.Lng, Lng: loc.Lat}
+		edited, err = b.Edit(edited, *loc)
+		assert.NoError(t, err)
+		assert.NotNil(t, edited.Location)
+	})
+
 	t.Run("Notify()", func(t *testing.T) {
 		assert.Equal(t, ErrBadRecipient, b.Notify(nil, Typing))
 		assert.NoError(t, b.Notify(to, Typing))