From b163052f82f70d259ba53759bd3470286201b479 Mon Sep 17 00:00:00 2001
From: Syfaro <syfaro@huefox.com>
Date: Sat, 20 Feb 2021 13:49:00 -0500
Subject: [PATCH] Handle InputMedia{Document,Audio} in media groups.

---
 bot_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++-
 configs.go  | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 helpers.go  |  2 +-
 types.go    | 12 ++++++++----
 4 files changed, 110 insertions(+), 6 deletions(-)

diff --git a/bot_test.go b/bot_test.go
index 537dcc3..af1e237 100644
--- a/bot_test.go
+++ b/bot_test.go
@@ -532,7 +532,7 @@ func TestSetWebhookWithoutCert(t *testing.T) {
 	bot.Request(RemoveWebhookConfig{})
 }
 
-func TestSendWithMediaGroup(t *testing.T) {
+func TestSendWithMediaGroupPhotoVideo(t *testing.T) {
 	bot, _ := getBot(t)
 
 	cfg := NewMediaGroup(ChatID, []interface{}{
@@ -555,6 +555,50 @@ func TestSendWithMediaGroup(t *testing.T) {
 	}
 }
 
+func TestSendWithMediaGroupDocument(t *testing.T) {
+	bot, _ := getBot(t)
+
+	cfg := NewMediaGroup(ChatID, []interface{}{
+		NewInputMediaDocument(FileURL("https://i.imgur.com/unQLJIb.jpg")),
+		NewInputMediaDocument("tests/image.jpg"),
+	})
+
+	messages, err := bot.SendMediaGroup(cfg)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if messages == nil {
+		t.Error("No received messages")
+	}
+
+	if len(messages) != len(cfg.Media) {
+		t.Errorf("Different number of messages: %d", len(messages))
+	}
+}
+
+func TestSendWithMediaGroupAudio(t *testing.T) {
+	bot, _ := getBot(t)
+
+	cfg := NewMediaGroup(ChatID, []interface{}{
+		NewInputMediaAudio("tests/audio.mp3"),
+		NewInputMediaAudio("tests/audio.mp3"),
+	})
+
+	messages, err := bot.SendMediaGroup(cfg)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if messages == nil {
+		t.Error("No received messages")
+	}
+
+	if len(messages) != len(cfg.Media) {
+		t.Errorf("Different number of messages: %d", len(messages))
+	}
+}
+
 func ExampleNewBotAPI() {
 	bot, err := NewBotAPI("MyAwesomeBotToken")
 	if err != nil {
diff --git a/configs.go b/configs.go
index 925766f..3231407 100644
--- a/configs.go
+++ b/configs.go
@@ -1805,6 +1805,30 @@ func prepareInputMediaParam(inputMedia interface{}, idx int) interface{} {
 			m.Thumb = fmt.Sprintf("attach://file-%d-thumb", idx)
 		}
 
+		return m
+	case InputMediaAudio:
+		switch m.Media.(type) {
+		case string, FileBytes, FileReader:
+			m.Media = fmt.Sprintf("attach://file-%d", idx)
+		}
+
+		switch m.Thumb.(type) {
+		case string, FileBytes, FileReader:
+			m.Thumb = fmt.Sprintf("attach://file-%d-thumb", idx)
+		}
+
+		return m
+	case InputMediaDocument:
+		switch m.Media.(type) {
+		case string, FileBytes, FileReader:
+			m.Media = fmt.Sprintf("attach://file-%d", idx)
+		}
+
+		switch m.Thumb.(type) {
+		case string, FileBytes, FileReader:
+			m.Thumb = fmt.Sprintf("attach://file-%d-thumb", idx)
+		}
+
 		return m
 	}
 
@@ -1847,6 +1871,38 @@ func prepareInputMediaFile(inputMedia interface{}, idx int) []RequestFile {
 				File: f,
 			})
 		}
+	case InputMediaDocument:
+		switch f := m.Media.(type) {
+		case string, FileBytes, FileReader:
+			files = append(files, RequestFile{
+				Name: fmt.Sprintf("file-%d", idx),
+				File: f,
+			})
+		}
+
+		switch f := m.Thumb.(type) {
+		case string, FileBytes, FileReader:
+			files = append(files, RequestFile{
+				Name: fmt.Sprintf("file-%d", idx),
+				File: f,
+			})
+		}
+	case InputMediaAudio:
+		switch f := m.Media.(type) {
+		case string, FileBytes, FileReader:
+			files = append(files, RequestFile{
+				Name: fmt.Sprintf("file-%d", idx),
+				File: f,
+			})
+		}
+
+		switch f := m.Thumb.(type) {
+		case string, FileBytes, FileReader:
+			files = append(files, RequestFile{
+				Name: fmt.Sprintf("file-%d", idx),
+				File: f,
+			})
+		}
 	}
 
 	return files
diff --git a/helpers.go b/helpers.go
index 8557970..87b5d4a 100644
--- a/helpers.go
+++ b/helpers.go
@@ -201,7 +201,7 @@ func NewInputMediaAudio(media interface{}) InputMediaAudio {
 }
 
 // NewInputMediaDocument creates a new InputMediaDocument.
-func NewInputMediaDocument(media string) InputMediaDocument {
+func NewInputMediaDocument(media interface{}) InputMediaDocument {
 	return InputMediaDocument{
 		BaseInputMedia: BaseInputMedia{
 			Type:  "document",
diff --git a/types.go b/types.go
index 2eb3ad5..bba539d 100644
--- a/types.go
+++ b/types.go
@@ -1112,10 +1112,11 @@ type BotCommand struct {
 
 // BaseInputMedia is a base type for the InputMedia types.
 type BaseInputMedia struct {
-	Type      string      `json:"type"`
-	Media     interface{} `json:"media"`
-	Caption   string      `json:"caption,omitempty"`
-	ParseMode string      `json:"parse_mode,omitempty"`
+	Type            string          `json:"type"`
+	Media           interface{}     `json:"media"`
+	Caption         string          `json:"caption,omitempty"`
+	CaptionEntities []MessageEntity `json:"caption_entities,omitempty"`
+	ParseMode       string          `json:"parse_mode,omitempty"`
 }
 
 // InputMediaPhoto is a photo to send as part of a media group.
@@ -1144,6 +1145,7 @@ type InputMediaAnimation struct {
 // InputMediaAudio is a audio to send as part of a media group.
 type InputMediaAudio struct {
 	BaseInputMedia
+	Thumb     interface{}
 	Duration  int    `json:"duration"`
 	Performer string `json:"performer"`
 	Title     string `json:"title"`
@@ -1152,6 +1154,8 @@ type InputMediaAudio struct {
 // InputMediaDocument is a audio to send as part of a media group.
 type InputMediaDocument struct {
 	BaseInputMedia
+	Thumb                       interface{}
+	DisableContentTypeDetection bool `json:"disable_content_type_detection,omitempty"`
 }
 
 // Error is an error containing extra information returned by the Telegram API.