git » chasquid » commit abf91ea

diff --git a/test/util/mail_diff b/test/util/mail_diff
index 761073d..26750e5 100755
--- a/test/util/mail_diff
+++ b/test/util/mail_diff
@@ -2,31 +2,10 @@
 
 import difflib
 import email.parser
+import itertools
 import mailbox
 import sys
 
-f1, f2 = sys.argv[1:3]
-
-expected = email.parser.Parser().parse(open(f1))
-
-mbox = mailbox.mbox(f2, create=False)
-msg = mbox[0]
-
-diff = False
-
-for h, val in expected.items():
-	if h not in msg:
-		print("Header missing: %r" % h)
-		diff = True
-		continue
-
-	if expected[h] == '*':
-		continue
-
-	if msg[h] != val:
-		print("Header %r differs: %r != %r" % (h, val, msg[h]))
-		diff = True
-
 
 def flexible_eq(expected, got):
     """Compare two strings, supporting wildcards.
@@ -48,7 +27,7 @@ def flexible_eq(expected, got):
             posG += 1
             continue
         if c == '*':
-            while got[posG] != '\n':
+            while posG < len(got) and got[posG] != '\n':
                 posG += 1
                 continue
             continue
@@ -58,20 +37,61 @@ def flexible_eq(expected, got):
 
         posG += 1
 
+    if posG != len(got):
+        # We got more than we expected.
+        return False
+
     return True
 
 
-if not flexible_eq(expected.get_payload(), msg.get_payload()):
-	diff = True
+def msg_equals(expected, msg):
+	"""Compare two messages recursively, using flexible_eq()."""
+	diff = False
+	for h, val in expected.items():
+		if h not in msg:
+			print("Header missing: %r" % h)
+			diff = True
+			continue
+
+		if expected[h] == '*':
+			continue
+
+		if not flexible_eq(val, msg[h]):
+			print("Header %r differs:" % h)
+			print("Exp: %r" % val)
+			print("Got: %r" % msg[h])
+			diff = True
+
+	if diff:
+		return False
 
 	if expected.is_multipart() != msg.is_multipart():
 		print("Multipart differs, expected %s, got %s" % (
 			expected.is_multipart(), msg.is_multipart()))
-	elif not msg.is_multipart():
-		exp = expected.get_payload().splitlines()
-		got = msg.get_payload().splitlines()
-		print("Payload differs:")
-		for l in difflib.ndiff(exp, got):
-			print(l)
-
-sys.exit(0 if not diff else 1)
+		return False
+
+	if expected.is_multipart():
+		for exp, got in itertools.izip_longest(expected.get_payload(), msg.get_payload()):
+			if not msg_equals(exp, got):
+				return False
+	else:
+		if not flexible_eq(expected.get_payload(), msg.get_payload()):
+			exp = expected.get_payload().splitlines()
+			got = msg.get_payload().splitlines()
+			print("Payload differs:")
+			for l in difflib.ndiff(exp, got):
+				print(l)
+			return False
+
+	return True
+
+
+if __name__ == "__main__":
+	f1, f2 = sys.argv[1:3]
+
+	expected = email.parser.Parser().parse(open(f1))
+
+	mbox = mailbox.mbox(f2, create=False)
+	msg = mbox[0]
+
+	sys.exit(0 if msg_equals(expected, msg) else 1)