Files
mautrix-telegram/pkg/gotd/telegram/query/photos/iter_test.go
T
2025-06-27 20:03:37 -07:00

85 lines
2.0 KiB
Go

package photos
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"go.mau.fi/mautrix-telegram/pkg/gotd/tg"
"go.mau.fi/mautrix-telegram/pkg/gotd/tgmock"
)
func generatePhotos(count int) []tg.PhotoClass {
r := make([]tg.PhotoClass, 0, count)
for i := 0; i < count; i++ {
r = append(r, &tg.Photo{
ID: int64(i + 1),
AccessHash: int64(i + 1),
FileReference: []uint8{uint8(i)},
})
}
return r
}
func result(r []tg.PhotoClass, count int) tg.PhotosPhotosClass {
return &tg.PhotosPhotosSlice{
Photos: r,
Count: count,
}
}
func TestIterator(t *testing.T) {
ctx := context.Background()
mock := tgmock.NewRequire(t)
limit := 10
totalRecords := 3 * limit
expected := generatePhotos(totalRecords)
raw := tg.NewClient(mock)
mock.ExpectCall(&tg.PhotosGetUserPhotosRequest{
UserID: &tg.InputUserSelf{},
Offset: 0,
Limit: limit,
}).ThenResult(result(expected[0:limit], totalRecords))
mock.ExpectCall(&tg.PhotosGetUserPhotosRequest{
UserID: &tg.InputUserSelf{},
Offset: limit,
Limit: limit,
}).ThenResult(result(expected[limit:2*limit], totalRecords))
mock.ExpectCall(&tg.PhotosGetUserPhotosRequest{
UserID: &tg.InputUserSelf{},
Offset: 2 * limit,
Limit: limit,
}).ThenResult(result(expected[2*limit:3*limit], totalRecords))
mock.ExpectCall(&tg.PhotosGetUserPhotosRequest{
UserID: &tg.InputUserSelf{},
Offset: 3 * limit,
Limit: limit,
}).ThenResult(result(expected[3*limit:], totalRecords))
iter := NewQueryBuilder(raw).GetUserPhotos(&tg.InputUserSelf{}).BatchSize(10).Iter()
i := 0
for iter.Next(ctx) {
require.Equal(t, expected[i], iter.Value().Photo)
i++
}
require.NoError(t, iter.Err())
require.Equal(t, totalRecords, i)
total, err := iter.Total(ctx)
require.NoError(t, err)
require.Equal(t, totalRecords, total)
mock.ExpectCall(&tg.PhotosGetUserPhotosRequest{
UserID: &tg.InputUserSelf{},
Offset: 0,
Limit: 1,
}).ThenResult(result(expected[:0], totalRecords))
total, err = iter.FetchTotal(ctx)
require.NoError(t, err)
require.Equal(t, totalRecords, total)
}