Keep the connection of a new WhatsApp link

[onechat] / whatsapp_fake_helper_test.go

diff --git a/whatsapp_fake_helper_test.go b/whatsapp_fake_helper_test.go
index e463cd094605600ec6a11497baced15929f03a38..c9bc8357efc79e527fb8ffe7f746f57b5aec1c77 100644 (file)
--- a/whatsapp_fake_helper_test.go
+++ b/whatsapp_fake_helper_test.go
@@ -67,28 +67,33 @@ var (
 
 // WhatsAppFake is a fake WhatsApp client. It keeps the event handler that the
 // daemon gives it, and Emit gives events to the handler, as the library does.
-// Each connection sends to connects, and each pairing sends its channel to
-// pairings, so that the test gives the items of the pairing. A connection fails
-// with connectErr, a logout with logoutErr, and a send with sendErr. groups are
-// the groups of the account. A send that does not fail gets the time
-// whatsappTestSendTime.
+// Each connection keeps its context and sends to connects, and each pairing
+// sends its channel to pairings, so that the test gives the items of the
+// pairing. A connection fails with connectErr, a logout with logoutErr, and a
+// send with sendErr. groups are the groups of the account. A send that does not
+// fail gets the time whatsappTestSendTime.
 type WhatsAppFake struct {
        connects chan struct{}
        pairings chan chan whatsmeow.QRChannelItem
 
-       mutex      sync.Mutex
-       handle     func(event any) bool
-       connectErr error
-       logoutErr  error
-       groups     []*types.GroupInfo
-       sendErr    error
-       phones     []string
-       logouts    int
-       sent       []WhatsAppFakeSend
+       mutex       sync.Mutex
+       handle      func(event any) bool
+       connectErr  error
+       logoutErr   error
+       groups      []*types.GroupInfo
+       sendErr     error
+       phones      []string
+       logouts     int
+       sent        []WhatsAppFakeSend
+       connections []context.Context
 }
 
-// ConnectContext tells the test that the daemon connected.
+// ConnectContext keeps ctx and tells the test that the daemon connected.
 func (f *WhatsAppFake) ConnectContext(ctx context.Context) error {
+       f.mutex.Lock()
+       f.connections = append(f.connections, ctx)
+       f.mutex.Unlock()
+
        f.connects <- struct{}{}
 
        f.mutex.Lock()
@@ -97,6 +102,14 @@ func (f *WhatsAppFake) ConnectContext(ctx context.Context) error {
        return f.connectErr
 }
 
+// Connections returns the contexts of the connections, in order.
+func (f *WhatsAppFake) Connections() []context.Context {
+       f.mutex.Lock()
+       defer f.mutex.Unlock()
+
+       return slices.Clone(f.connections)
+}
+
 // Disconnect does nothing.
 func (f *WhatsAppFake) Disconnect() {
 }