X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/f2e67d5579593e2d4f1cb2d9165b851c52a06715..a85d214071706753f358f0af23cfd525c28dbfe0:/whatsapp_fake_helper_test.go

diff --git a/whatsapp_fake_helper_test.go b/whatsapp_fake_helper_test.go
index e463cd0..c9bc835 100644
--- 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() {
 }