From cc0af4e7906e02bd0aca6424c73960118e6dc2cc Mon Sep 17 00:00:00 2001
From: Valentin Lorentz <progval+git@progval.net>
Date: Wed, 6 May 2020 18:29:17 +0200
Subject: [PATCH] Include start and end messages in a batch.

They may have useful server tags (especially when we start
supporting labeled-response).
---
 src/irclib.py       |  3 ++-
 test/test_irclib.py | 14 ++++++++------
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/irclib.py b/src/irclib.py
index a8048f775..85c01f3b4 100644
--- a/src/irclib.py
+++ b/src/irclib.py
@@ -845,9 +845,10 @@ class IrcState(IrcCommandDispatcher, log.Firewalled):
             batch_type = msg.args[1]
             batch_arguments = tuple(msg.args[2:])
             self.batches[batch_name] = Batch(type=batch_type,
-                    arguments=batch_arguments, messages=[])
+                    arguments=batch_arguments, messages=[msg])
         elif msg.args[0].startswith('-'):
             batch = self.batches.pop(batch_name)
+            batch.messages.append(msg)
             msg.tag('batch', batch)
         else:
             assert False, msg.args[0]
diff --git a/test/test_irclib.py b/test/test_irclib.py
index cdc0ea62d..8d56a8728 100644
--- a/test/test_irclib.py
+++ b/test/test_irclib.py
@@ -748,12 +748,13 @@ class IrcTestCase(SupyTestCase):
     def testBatch(self):
         self.irc.reset()
         self.irc.feedMsg(ircmsgs.IrcMsg(':someuser1 JOIN #foo'))
-        self.irc.feedMsg(ircmsgs.IrcMsg(':host BATCH +name netjoin'))
-        m1 = ircmsgs.IrcMsg('@batch=name :someuser2 JOIN #foo')
+        m1 = ircmsgs.IrcMsg(':host BATCH +name netjoin')
         self.irc.feedMsg(m1)
-        self.irc.feedMsg(ircmsgs.IrcMsg(':someuser3 JOIN #foo'))
-        m2 = ircmsgs.IrcMsg('@batch=name :someuser4 JOIN #foo')
+        m2 = ircmsgs.IrcMsg('@batch=name :someuser2 JOIN #foo')
         self.irc.feedMsg(m2)
+        self.irc.feedMsg(ircmsgs.IrcMsg(':someuser3 JOIN #foo'))
+        m3 = ircmsgs.IrcMsg('@batch=name :someuser4 JOIN #foo')
+        self.irc.feedMsg(m3)
         class Callback(irclib.IrcCallback):
             batch = None
             def name(self):
@@ -763,10 +764,11 @@ class IrcTestCase(SupyTestCase):
         c = Callback()
         self.irc.addCallback(c)
         try:
-            self.irc.feedMsg(ircmsgs.IrcMsg(':host BATCH -name'))
+            m4 = ircmsgs.IrcMsg(':host BATCH -name')
+            self.irc.feedMsg(m4)
         finally:
             self.irc.removeCallback(c.name())
-        self.assertEqual(c.batch, irclib.Batch('netjoin', (), [m1, m2]))
+        self.assertEqual(c.batch, irclib.Batch('netjoin', (), [m1, m2, m3, m4]))
 
 class SaslTestCase(SupyTestCase):
     def setUp(self):