allow to pass file as argument

1 files changed, 13 insertions, 3 deletions

diff --git a/src/leap/mail/imap/tests/walktree.py b/src/leap/mail/imap/tests/walktree.py
index 1626f65..f3cbcb0 100644
--- a/src/leap/mail/imap/tests/walktree.py
+++ b/src/leap/mail/imap/tests/walktree.py
@@ -18,12 +18,14 @@
 Tests for the walktree module.
 """
 import os
+import sys
 from email import parser
 
 from leap.mail import walk as W
 
 DEBUG = os.environ.get("BITMASK_MAIL_DEBUG")
 
+
 p = parser.Parser()
 
 # TODO pass an argument of the type of message
@@ -31,9 +33,17 @@ p = parser.Parser()
 ##################################################
 # Input from hell
 
-#msg = p.parse(open('rfc822.multi-signed.message'))
-#msg = p.parse(open('rfc822.plain.message'))
-msg = p.parse(open('rfc822.multi-minimal.message'))
+if len(sys.argv) > 1:
+    FILENAME = sys.argv[1]
+else:
+    FILENAME = "rfc822.multi-minimal.message"
+
+"""
+FILENAME = "rfc822.multi-signed.message"
+FILENAME = "rfc822.plain.message"
+"""
+
+msg = p.parse(open(FILENAME))
 DO_CHECK = False
 #################################################