test(e2etest): fix

telegram/internal/e2etest/echo_bot.go

@@ -5,7 +5,6 @@ import (
 	"strconv"
 	"sync"
 
-	"github.com/cenkalti/backoff/v4"
 	"github.com/go-faster/errors"
 	"go.uber.org/zap"
 
@@ -71,39 +70,26 @@ func (b EchoBot) login(ctx context.Context, client *telegram.Client) (*tg.User,
 	}
 
 	var me *tg.User
-	if err := retryFloodWait(ctx, func() (err error) {
+	if err := retry(ctx, func() (err error) {
 		me, err = client.Self(ctx)
 		return err
 	}); err != nil {
 		return nil, err
 	}
 
 	expectedUsername := "echobot" + strconv.FormatInt(me.ID, 10)
-	raw := tg.NewClient(waitInvoker{prev: client})
+	raw := tg.NewClient(retryInvoker{prev: client})
 	_, err := raw.AccountUpdateUsername(ctx, expectedUsername)
 	if err != nil {
 		if !tgerr.Is(err, tg.ErrUsernameNotModified) {
 			return nil, errors.Wrap(err, "update username")
 		}
 	}
-
-	if err := backoff.Retry(func() error {
-		me, err = client.Self(ctx)
-		if err != nil {
-			if ok, err := tgerr.FloodWait(ctx, err); ok {
-				return err
-			}
-
-			return backoff.Permanent(errors.Wrap(err, "get self"))
-		}
-
-		if me.Username != expectedUsername {
-			return errors.Errorf("expected username %q, got %q", expectedUsername, me.Username)
-		}
-
-		return nil
-	}, backoff.WithContext(backoff.NewExponentialBackOff(), ctx)); err != nil {
-		return nil, err
+	me, err = retryResult(ctx, func() (*tg.User, error) {
+		return client.Self(ctx)
+	})
+	if me.Username != expectedUsername {
+		return nil, errors.Errorf("expected username %q, got %q", expectedUsername, me.Username)
 	}
 
 	return me, nil
@@ -125,14 +111,19 @@ func (b EchoBot) handler(client *telegram.Client) tg.NewMessageHandler {
 		}
 
 		if dialogsUsers.empty() {
-			dialogs, err := raw.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
-				Limit:      100,
-				OffsetPeer: &tg.InputPeerEmpty{},
+			dialogs, err := retryResult(ctx, func() (tg.MessagesDialogsClass, error) {
+				dialogs, err := raw.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
+					Limit:      100,
+					OffsetPeer: &tg.InputPeerEmpty{},
+				})
+				if err != nil {
+					return nil, errors.Wrap(err, "get dialogs")
+				}
+				return dialogs, nil
 			})
 			if err != nil {
 				return errors.Wrap(err, "get dialogs")
 			}
-
 			if dlg, ok := dialogs.AsModified(); ok {
 				dialogsUsers.add(dlg.GetUsers()...)
 			}
@@ -154,7 +145,10 @@ func (b EchoBot) handler(client *telegram.Client) tg.NewMessageHandler {
 					zap.String("username", user.Username),
 				)
 
-				if _, err := sender.To(user.AsInputPeer()).Text(ctx, m.Message); err != nil {
+				if err := retry(ctx, func() error {
+					_, err := sender.To(user.AsInputPeer()).Text(ctx, m.Message)
+					return err
+				}); err != nil {
 					return errors.Wrap(err, "send message")
 				}
 				return nil

telegram/internal/e2etest/retry.go

@@ -0,0 +1,61 @@
+package e2etest
+
+import (
+	"context"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+
+	"github.com/gotd/td/bin"
+	"github.com/gotd/td/tg"
+	"github.com/gotd/td/tgerr"
+)
+
+type retryInvoker struct {
+	prev tg.Invoker
+}
+
+func retryResult[T any](ctx context.Context, cb func() (T, error)) (T, error) {
+	var zero T
+	return backoff.RetryWithData[T](func() (T, error) {
+		res, err := cb()
+		if err != nil {
+			if tgerr.IsCode(err, -500) {
+				return zero, err
+			}
+			if tgerr.Is(err, "CONNECTION_NOT_INITED") {
+				return zero, err
+			}
+			if ok, err := tgerr.FloodWait(ctx, err); ok {
+				return zero, err
+			}
+			return zero, backoff.Permanent(err)
+		}
+		return res, nil
+	}, backoff.WithContext(backoff.NewConstantBackOff(time.Millisecond*500), ctx))
+}
+
+func retry(ctx context.Context, cb func() error) error {
+	return backoff.Retry(func() error {
+		if err := cb(); err != nil {
+			if tgerr.IsCode(err, -500) {
+				return err
+			}
+			if tgerr.Is(err, "CONNECTION_NOT_INITED") {
+				return err
+			}
+			if ok, err := tgerr.FloodWait(ctx, err); ok {
+				return err
+			}
+			return backoff.Permanent(err)
+		}
+
+		return nil
+	}, backoff.WithContext(backoff.NewExponentialBackOff(), ctx))
+}
+
+func (w retryInvoker) Invoke(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+	return retry(ctx, func() error {
+		return w.prev.Invoke(ctx, input, output)
+	})
+}

telegram/internal/e2etest/terentyev.go

@@ -62,7 +62,7 @@ func (u User) Run(ctx context.Context) error {
 	dispatcher := tg.NewUpdateDispatcher()
 	dispatcher.OnNewMessage(u.messageHandler)
 	client := u.suite.Client(u.logger, dispatcher)
-	sender := message.NewSender(tg.NewClient(waitInvoker{prev: client}))
+	sender := message.NewSender(tg.NewClient(retryInvoker{prev: client}))
 
 	return client.Run(ctx, func(ctx context.Context) error {
 		if err := u.suite.RetryAuthenticate(ctx, client.Auth()); err != nil {