diff --git a/tests/scripts/audit-validity-dates.py b/tests/scripts/audit-validity-dates.py
index d74c6f8..09559dc 100755
--- a/tests/scripts/audit-validity-dates.py
+++ b/tests/scripts/audit-validity-dates.py
@@ -42,15 +42,20 @@
 from generate_test_code import parse_test_data as parse_suite_data
 from generate_test_code import FileWrapper
 
+import scripts_path # pylint: disable=unused-import
+from mbedtls_dev import build_tree
+
 class DataType(Enum):
     CRT = 1 # Certificate
     CRL = 2 # Certificate Revocation List
     CSR = 3 # Certificate Signing Request
 
+
 class DataFormat(Enum):
     PEM = 1 # Privacy-Enhanced Mail
     DER = 2 # Distinguished Encoding Rules
 
+
 class AuditData:
     """Store data location, type and validity period of X.509 objects."""
     #pylint: disable=too-few-public-methods
@@ -78,6 +83,7 @@
         else:
             raise ValueError("Unsupported file_type: {}".format(self.data_type))
 
+
 class X509Parser:
     """A parser class to parse crt/crl/csr file or data in PEM/DER format."""
     PEM_REGEX = br'-{5}BEGIN (?P<type>.*?)-{5}\n(?P<data>.*?)-{5}END (?P=type)-{5}\n'
@@ -167,6 +173,7 @@
             return False
         return True
 
+
 class Auditor:
     """A base class for audit."""
     def __init__(self, logger):
@@ -231,15 +238,8 @@
     @staticmethod
     def find_test_dir():
         """Get the relative path for the MbedTLS test directory."""
-        if os.path.isdir('tests'):
-            tests_dir = 'tests'
-        elif os.path.isdir('suites'):
-            tests_dir = '.'
-        elif os.path.isdir('../suites'):
-            tests_dir = '..'
-        else:
-            raise Exception("Mbed TLS source tree not found")
-        return tests_dir
+        return os.path.relpath(build_tree.guess_mbedtls_root() + '/tests')
+
 
 class TestDataAuditor(Auditor):
     """Class for auditing files in tests/data_files/"""
@@ -255,6 +255,7 @@
                       if os.path.isfile(f)]
         return data_files
 
+
 class SuiteDataAuditor(Auditor):
     """Class for auditing files in tests/suites/*.data"""
     def __init__(self, options):
@@ -294,6 +295,7 @@
 
         return audit_data_list
 
+
 def list_all(audit_data: AuditData):
     print("{}\t{}\t{}\t{}".format(
         audit_data.not_valid_before.isoformat(timespec='seconds'),