Make main() suitable to being called from python

Don't call sys.exit(), and don't clobber the working directory.

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

tests/scripts/test_psa_compliance.py

@@ -18,7 +18,7 @@ EXPECTED_FAILURES = {
 PSA_ARCH_TESTS_REPO = 'https://github.com/ronald-cron-arm/psa-arch-tests.git'
 PSA_ARCH_TESTS_REF = 'crypto1.0-3.0'
 
-#pylint: disable=too-many-statements
+#pylint: disable=too-many-branches,too-many-statements
 def main():
     mbedtls_dir = os.getcwd()
 
@@ -30,6 +30,7 @@ def main():
         os.mkdir(psa_arch_tests_dir)
     except FileExistsError:
         pass
+    try:
         os.chdir(psa_arch_tests_dir)
 
         subprocess.check_call(['git', 'init'])
@@ -46,7 +47,8 @@ def main():
 
         #pylint: disable=bad-continuation
         subprocess.check_call([
-        'cmake', '..', '-GUnix Makefiles',
+            'cmake', '..',
+                     '-GUnix Makefiles',
                      '-DTARGET=tgt_dev_apis_stdc',
                      '-DTOOLCHAIN=HOST_GCC',
                      '-DSUITE=CRYPTO',
@@ -100,11 +102,13 @@ def main():
                       'EXPECTED_FAILURES in tests/scripts/compliance_test.py')
                 print()
             print('FAILED')
-        sys.exit(1)
+            return 1
         else:
-        os.chdir(mbedtls_dir)
             shutil.rmtree(psa_arch_tests_dir)
             print('SUCCESS')
+            return 0
+    finally:
+        os.chdir(mbedtls_dir)
 
 if __name__ == '__main__':
-    main()
+    sys.exit(main())