From 260a67b811cc66e4ce098296a7a4b3f570cb7a0b Mon Sep 17 00:00:00 2001
From: Nikita <nenten137@gmail.com>
Date: Fri, 1 Mar 2024 21:58:56 +0200
Subject: [PATCH] media: implement SetCaption in Album

---
 media.go      | 18 ++++++++++++++++++
 media_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 media_test.go

diff --git a/media.go b/media.go
index 9ef1127..bdce555 100644
--- a/media.go
+++ b/media.go
@@ -45,6 +45,24 @@ type Inputtable interface {
 // Album lets you group multiple media into a single message.
 type Album []Inputtable
 
+func (a Album) SetCaption(caption string) {
+	if len(a) < 1 {
+		return
+	}
+	switch a[0].MediaType() {
+	case "audio":
+		a[0].(*Audio).Caption = caption
+	case "video":
+		a[0].(*Video).Caption = caption
+	case "document":
+		a[0].(*Document).Caption = caption
+	case "photo":
+		a[0].(*Photo).Caption = caption
+	case "animation":
+		a[0].(*Animation).Caption = caption
+	}
+}
+
 // Photo object represents a single photo file.
 type Photo struct {
 	File
diff --git a/media_test.go b/media_test.go
new file mode 100644
index 0000000..2f545e0
--- /dev/null
+++ b/media_test.go
@@ -0,0 +1,43 @@
+package telebot
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestAlbumSetCaption(t *testing.T) {
+	var a Album
+	a = append(a, &Photo{Caption: "wrong_caption"})
+	a = append(a, &Photo{Caption: "t"})
+	a.SetCaption("correct_caption")
+	assert.Equal(t, "correct_caption", a[0].InputMedia().Caption)
+	assert.Equal(t, "t", a[1].InputMedia().Caption)
+
+	a = Album{}
+	a = append(a, &Animation{Caption: "wrong_caption"})
+	a = append(a, &Photo{Caption: "t"})
+	a.SetCaption("correct_caption")
+	assert.Equal(t, "correct_caption", a[0].InputMedia().Caption)
+	assert.Equal(t, "t", a[1].InputMedia().Caption)
+
+	a = Album{}
+	a = append(a, &Audio{Caption: "wrong_caption"})
+	a = append(a, &Photo{Caption: "t"})
+	a.SetCaption("correct_caption")
+	assert.Equal(t, "correct_caption", a[0].InputMedia().Caption)
+	assert.Equal(t, "t", a[1].InputMedia().Caption)
+
+	a = Album{}
+	a = append(a, &Document{Caption: "wrong_caption"})
+	a = append(a, &Photo{Caption: "t"})
+	a.SetCaption("correct_caption")
+	assert.Equal(t, "correct_caption", a[0].InputMedia().Caption)
+	assert.Equal(t, "t", a[1].InputMedia().Caption)
+
+	a = Album{}
+	a = append(a, &Video{Caption: "wrong_caption"})
+	a = append(a, &Photo{Caption: "t"})
+	a.SetCaption("correct_caption")
+	assert.Equal(t, "correct_caption", a[0].InputMedia().Caption)
+	assert.Equal(t, "t", a[1].InputMedia().Caption)
+}