diff --git a/telegram/client_external_test.go b/telegram/client_external_test.go index 7fa37cc6ab..530fa17fa0 100644 --- a/telegram/client_external_test.go +++ b/telegram/client_external_test.go @@ -1,25 +1,30 @@ package telegram_test import ( + "bytes" "context" "crypto/rand" "strings" "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/zap" "go.uber.org/zap/zapcore" "go.uber.org/zap/zaptest" "golang.org/x/xerrors" + "github.com/gotd/td/internal/crypto" "github.com/gotd/td/internal/tdsync" "github.com/gotd/td/internal/testutil" "github.com/gotd/td/session" "github.com/gotd/td/telegram" "github.com/gotd/td/telegram/dcs" + "github.com/gotd/td/telegram/downloader" "github.com/gotd/td/telegram/internal/e2etest" "github.com/gotd/td/telegram/internal/tgtest" + "github.com/gotd/td/telegram/message" "github.com/gotd/td/tg" "github.com/gotd/td/transport" ) @@ -103,3 +108,111 @@ func TestExternalE2EUsersDialog(t *testing.T) { require.NoError(t, g.Wait()) } + +func TestExternalE2EPartialUpload(t *testing.T) { + testutil.SkipExternal(t) + + // Testing partial uploads. + // + // Partial uploads can ease some invariants: + // * Know total chunk count only on last chunk + // * Upload first chunk (part=0) after last chunk + // + // Note that overwriting (re-uploading part with different content) is + // unreliable and works only on very small files, so if you want to + // finalize first chunk later, don't upload it until finalization. + + log := zaptest.NewLogger(t).WithOptions(zap.IncreaseLevel(zapcore.InfoLevel)) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + cfg := e2etest.TestConfig{ + AppID: telegram.TestAppID, + AppHash: telegram.TestAppHash, + DcID: 2, + } + suite := e2etest.NewSuite(tgtest.NewSuite(ctx, t, log), cfg, rand.Reader) + + client := suite.Client(log, nil) + require.NoError(t, client.Run(ctx, func(ctx context.Context) error { + if err := suite.RetryAuthenticate(ctx, client); err != nil { + return err + } + + raw := tg.NewClient(client) + fileID, err := crypto.RandInt64(rand.Reader) + if err != nil { + return err + } + + const ( + totalParts = 10 + allocParts = totalParts + 1 + lastPart = totalParts - 1 + chunkSize = 1024 * 512 / 4 + lastChunkSize = chunkSize / 2 + ) + + // Uploading chunks from 1 to totalParts. + var totalSize int + for i := 1; i < totalParts; i++ { + req := &tg.UploadSaveBigFilePartRequest{ + FileID: fileID, + FilePart: i, + FileTotalParts: allocParts, + Bytes: bytes.Repeat([]byte{byte(i + 1)}, chunkSize), + } + + // Last chunk should finalize part size, changing it from allocated parts + // to real part count. + if req.FilePart == lastPart { + req.Bytes = req.Bytes[:lastChunkSize] + req.FileTotalParts = totalParts + } + totalSize += len(req.Bytes) + if _, err := raw.UploadSaveBigFilePart(ctx, req); err != nil { + return xerrors.Errorf("part: %w", err) + } + t.Logf("Uploaded part %d of %d (len: %d)", req.FilePart+1, req.FileTotalParts, len(req.Bytes)) + } + + // Uploading first chunk. + req := &tg.UploadSaveBigFilePartRequest{ + FileID: fileID, + FilePart: 0, + FileTotalParts: totalParts, + Bytes: bytes.Repeat([]byte{50}, chunkSize), + } + totalSize += len(req.Bytes) + if _, err := raw.UploadSaveBigFilePart(ctx, req); err != nil { + return xerrors.Errorf("part: %w", err) + } + + f := &tg.InputFileBig{ + ID: fileID, + Parts: totalParts, + Name: "data.bin", + } + + mc, err := message.NewSender(raw).Self().UploadMedia(ctx, message.File(f)) + if err != nil { + return xerrors.Errorf("file: %w", err) + } + + doc, ok := mc.(*tg.MessageMediaDocument).Document.AsNotEmpty() + if !ok { + return xerrors.New("bad doc") + } + + buf := new(bytes.Buffer) + if _, err := downloader.NewDownloader().Download(raw, doc.AsInputDocumentFileLocation()).Stream(ctx, buf); err != nil { + return err + } + + assert.Equal(t, totalSize, buf.Len()) + assert.Equal(t, req.Bytes, buf.Bytes()[:chunkSize]) + + t.Logf("File uploaded (id=%d, parts=%d, name=%s, size %d kb)", f.ID, f.Parts, f.Name, totalSize/1024) + + return nil + })) +}