3 files changed, 48 insertions, 20 deletions

diff --git a/test/clitest.py b/test/clitest.py
index 0a29c02..c913bdf 100644
--- a/test/clitest.py
+++ b/test/clitest.py
@@ -11,9 +11,15 @@ import subprocess
 import sys
 import tarfile
 
-sys.path.append('..')
-from libmat import mat
 import test
+MAT_PATH = 'mat'
+if test.IS_LOCAL is True:
+    # Are we testing the _local_ version of MAT?
+    sys.path.insert(0, '..')
+    MAT_PATH = '../mat'
+# else it will be in the path
+
+from libmat import mat
 
 
 class TestRemovecli(test.MATTest):
@@ -24,14 +30,14 @@ class TestRemovecli(test.MATTest):
     def test_remove(self):
         """make sure that the cli remove all compromizing meta"""
         for _, dirty in self.file_list:
-            subprocess.call(['../mat', '--add2archive', dirty])
+            subprocess.call([MAT_PATH, '--add2archive', dirty])
             current_file = mat.create_class_file(dirty, False, add2archive=True, low_pdf_quality=True)
             self.assertTrue(current_file.is_clean())
 
     def test_remove_empty(self):
         """Test removal with clean files\n"""
         for clean, _ in self.file_list:
-            subprocess.call(['../mat', '--add2archive', clean])
+            subprocess.call([MAT_PATH, '--add2archive', clean])
             current_file = mat.create_class_file(clean, False, add2archive=True, low_pdf_quality=True)
             self.assertTrue(current_file.is_clean())
 
@@ -44,7 +50,7 @@ class TestListcli(test.MATTest):
     def test_list_clean(self):
         """check if get_meta returns meta"""
         for clean, _ in self.file_list:
-            proc = subprocess.Popen(['../mat', '-d', clean],
+            proc = subprocess.Popen([MAT_PATH, '-d', clean],
                                     stdout=subprocess.PIPE)
             stdout, _ = proc.communicate()
             self.assertEqual(str(stdout).strip('\n'), "[+] File %s \
@@ -53,7 +59,7 @@ class TestListcli(test.MATTest):
     def test_list_dirty(self):
         """check if get_meta returns all the expected meta"""
         for _, dirty in self.file_list:
-            proc = subprocess.Popen(['../mat', '-d', dirty],
+            proc = subprocess.Popen([MAT_PATH, '-d', dirty],
                                     stdout=subprocess.PIPE)
             stdout, _ = proc.communicate()
             self.assertNotEqual(str(stdout), "[+] File %s :\n No\
@@ -68,7 +74,7 @@ class TestisCleancli(test.MATTest):
     def test_clean(self):
         """test is_clean on clean files"""
         for clean, _ in self.file_list:
-            proc = subprocess.Popen(['../mat', '-c', clean],
+            proc = subprocess.Popen([MAT_PATH, '-c', clean],
                                     stdout=subprocess.PIPE)
             stdout, _ = proc.communicate()
             self.assertEqual(str(stdout).strip('\n'), '[+] %s is clean' % clean)
@@ -76,7 +82,7 @@ class TestisCleancli(test.MATTest):
     def test_dirty(self):
         """test is_clean on dirty files"""
         for _, dirty in self.file_list:
-            proc = subprocess.Popen(['../mat', '-c', dirty],
+            proc = subprocess.Popen([MAT_PATH, '-c', dirty],
                                     stdout=subprocess.PIPE)
             stdout, _ = proc.communicate()
             self.assertEqual(str(stdout).strip('\n'), '[+] %s is not clean' % dirty)
@@ -89,21 +95,21 @@ class TestFileAttributes(unittest.TestCase):
 
     def test_not_writtable(self):
         """ test MAT's behaviour on non-writable file"""
-        proc = subprocess.Popen(['../mat', 'not_writtable'],
+        proc = subprocess.Popen([MAT_PATH, 'not_writtable'],
                                 stdout=subprocess.PIPE)
         stdout, _ = proc.communicate()
         self.assertEqual(str(stdout).strip('\n'), '[-] %s is not writable' % 'not_writtable')
 
     def test_not_exist(self):
         """ test MAT's behaviour on non-existent file"""
-        proc = subprocess.Popen(['../mat', 'ilikecookies'],
+        proc = subprocess.Popen([MAT_PATH, 'ilikecookies'],
                                 stdout=subprocess.PIPE)
         stdout, _ = proc.communicate()
         self.assertEqual(str(stdout).strip('\n'), 'Unable to process  %s' % 'ilikecookies')
 
     def test_empty(self):
         """ test MAT's behaviour on empty file"""
-        proc = subprocess.Popen(['../mat', 'empty_file'], stdout=subprocess.PIPE)
+        proc = subprocess.Popen([MAT_PATH, 'empty_file'], stdout=subprocess.PIPE)
         stdout, _ = proc.communicate()
         self.assertEqual(str(stdout).strip('\n'), 'Unable to process  %s' % 'ilikecookies')
 
@@ -114,13 +120,13 @@ class TestUnsupported(test.MATTest):
         """
         tarpath = os.path.join(self.tmpdir, "test.tar.bz2")
         tar = tarfile.open(tarpath, "w")
-        for f in ('../mat.desktop', '../README.security', '../setup.py'):
-            tar.add(f, f[3:])  # trim '../'
+        for f in ('libtest.py', 'test.py', 'clitest.py'):
+            tar.add(f, f)
         tar.close()
-        proc = subprocess.Popen(['../mat', tarpath], stdout=subprocess.PIPE)
+        proc = subprocess.Popen([MAT_PATH, tarpath], stdout=subprocess.PIPE)
         stdout, _ = proc.communicate()
         self.assertTrue('It contains unsupported filetypes:' \
-                        '\n- mat.desktop\n- README.security\n- setup.py\n'
+                        '\n- libtest.py\n- test.py\n- clitest.py\n'
                         in str(stdout))
 
 
diff --git a/test/libtest.py b/test/libtest.py
index 25d7426..03ae4a1 100644
--- a/test/libtest.py
+++ b/test/libtest.py
@@ -10,10 +10,14 @@ import sys
 import shutil
 import tarfile
 import tempfile
-import test
 import unittest
 
-sys.path.append('..')
+import test
+if test.IS_LOCAL is True:
+    # Are we testing the _local_ version of MAT?
+    sys.path.insert(0, '..')
+# else it will be in the path
+
 import libmat
 
 
@@ -155,12 +159,12 @@ class TestArchiveProcessing(test.MATTest):
         """
         tarpath = os.path.join(self.tmpdir, "test.tar.bz2")
         tar = tarfile.open(tarpath, "w")
-        for f in ('../mat.desktop', '../README.security', '../setup.py'):
-            tar.add(f, f[3:])  # trim '../'
+        for f in ('libtest.py', 'test.py', 'clitest.py'):
+            tar.add(f, f)
         tar.close()
         current_file = libmat.mat.create_class_file(tarpath, False, add2archive=False)
         unsupported_files = set(current_file.is_clean(list_unsupported=True))
-        self.assertEqual(unsupported_files, {'mat.desktop', 'README.security', 'setup.py'})
+        self.assertEqual(unsupported_files, {'libtest.py', 'test.py', 'clitest.py'})
 
     def test_archive_unwritable_content(self):
         path = os.path.join(self.tmpdir, './unwritable_content.zip')
diff --git a/test/test.py b/test/test.py
index 40cbf0c..e45912e 100644
--- a/test/test.py
+++ b/test/test.py
@@ -15,6 +15,8 @@ import sys
 import tempfile
 import unittest
 
+IS_LOCAL = True
+
 VERBOSITY = 15
 
 clean = glob.glob('clean*')
@@ -72,6 +74,22 @@ class MATTest(unittest.TestCase):
 if __name__ == '__main__':
     import clitest
     import libtest
+    import argparse
+
+    parser = argparse.ArgumentParser(description='MAT testsuite')
+    parser.add_argument('-l', '--local', action='store_true',
+            help='Test the local version of mat')
+    parser.add_argument('-s', '--system', action='store_true',
+            help='Test the system-wide version of mat')
+
+    if parser.parse_args().local is True:
+        IS_LOCAL = True
+    elif parser.parse_args().system is True:
+        IS_LOCAL = False
+    else:
+        print('Please specify either --local or --system')
+        sys.exit(1)
+
 
     SUITE = unittest.TestSuite()
     SUITE.addTests(clitest.get_tests())