diff --git a/tests/treble_sepolicy_tests.py b/tests/treble_sepolicy_tests.py
index 3d6a480..2c0cef3 100644
--- a/tests/treble_sepolicy_tests.py
+++ b/tests/treble_sepolicy_tests.py
@@ -1,7 +1,9 @@
 from optparse import OptionParser
 from optparse import Option, OptionValueError
 import os
+import mini_parser
 import policy
+from policy import MatchPathPrefix
 import re
 import sys
 
@@ -69,26 +71,13 @@
 appdomains = set()
 vendordomains = set()
 
-###
-# Check whether the regex will match a file path starting with the provided
-# prefix
-#
-# Compares regex entries in file_contexts with a path prefix. Regex entries
-# are often more specific than this file prefix. For example, the regex could
-# be /system/bin/foo\.sh and the prefix could be /system. This function
-# loops over the regex removing characters from the end until
-# 1) there is a match - return True or 2) run out of characters - return
-#    False.
-#
-def MatchPathPrefix(pathregex, prefix):
-    for i in range(len(pathregex), 0, -1):
-        try:
-            pattern = re.compile('^' + pathregex[0:i] + "$")
-        except:
-            continue
-        if pattern.match(prefix):
-            return True
-    return False
+# compat vars
+alltypes = set()
+oldalltypes = set()
+compatMapping = None
+
+# Distinguish between PRODUCT_FULL_TREBLE and PRODUCT_FULL_TREBLE_OVERRIDE
+FakeTreble = False
 
 def GetAllDomains(pol):
     global alldomains
@@ -107,12 +96,11 @@
             alldomains[d].appdomain = True
             appdomains.add(d)
 
-
 def GetCoreDomains():
     global alldomains
     global coredomains
     for d in alldomains:
-        # TestCoredomainViolators will verify if coredomain was incorrectly
+        # TestCoredomainViolations will verify if coredomain was incorrectly
         # applied.
         if "coredomain" in alldomains[d].attributes:
             alldomains[d].coredomain = True
@@ -144,7 +132,7 @@
 #
 def GetDomainEntrypoints(pol):
     global alldomains
-    for x in pol.QueryTERule(tclass="file", perms=["entrypoint"]):
+    for x in pol.QueryExpandedTERule(tclass=set(["file"]), perms=set(["entrypoint"])):
         if not x.sctx in alldomains:
             continue
         alldomains[x.sctx].entrypoints.append(str(x.tctx))
@@ -167,6 +155,12 @@
         for result in pol.QueryTypeAttribute(domain, False):
             alldomains[domain].attributes.add(result)
 
+def GetAllTypes(pol, oldpol):
+    global alltypes
+    global oldalltypes
+    alltypes = pol.GetAllTypes(False)
+    oldalltypes = oldpol.GetAllTypes(False)
+
 def setup(pol):
     GetAllDomains(pol)
     GetAttributes(pol)
@@ -174,6 +168,21 @@
     GetAppDomains()
     GetCoreDomains()
 
+# setup for the policy compatibility tests
+def compatSetup(pol, oldpol, mapping):
+    global compatMapping
+
+    GetAllTypes(pol, oldpol)
+    compatMapping = mapping
+
+def DomainsWithAttribute(attr):
+    global alldomains
+    domains = []
+    for domain in alldomains:
+        if attr in alldomains[domain].attributes:
+            domains.append(domain)
+    return domains
+
 #############################################################
 # Tests
 #############################################################
@@ -209,6 +218,75 @@
     return ret
 
 ###
+# Make sure that any new type introduced in the new policy that was not present
+# in the old policy has been recorded in the mapping file.
+def TestNoUnmappedNewTypes():
+    global alltypes
+    global oldalltypes
+    global compatMapping
+    newt = alltypes - oldalltypes
+    ret = ""
+    violators = []
+
+    for n in newt:
+        if compatMapping.rTypeattributesets.get(n) is None:
+            violators.append(n)
+
+    if len(violators) > 0:
+        ret += "SELinux: The following types were found added to the policy "
+        ret += "without an entry into the compatibility mapping file(s) found "
+        ret += "in private/compat/" + compatMapping.apiLevel + "/"
+        ret +=  compatMapping.apiLevel + "[.ignore].cil/n"
+        ret += " ".join(str(x) for x in sorted(violators)) + "\n"
+    return ret
+
+###
+# Make sure that any public type removed in the current policy has its
+# declaration added to the mapping file for use in non-platform policy
+def TestNoUnmappedRmTypes():
+    global alltypes
+    global oldalltypes
+    global compatMapping
+    rmt = oldalltypes - alltypes
+    ret = ""
+    violators = []
+
+    for o in rmt:
+        if o in compatMapping.pubtypes and not o in compatMapping.types:
+            violators.append(o)
+
+    if len(violators) > 0:
+        ret += "SELinux: The following formerly public types were removed from "
+        ret += "policy without a declaration in the compatibility mapping "
+        ret += "file(s) found in prebuilts/api/" + compatMapping.apiLevel + "/\n"
+        ret += " ".join(str(x) for x in sorted(violators)) + "\n"
+    return ret
+
+def TestTrebleCompatMapping():
+    ret = TestNoUnmappedNewTypes()
+    ret += TestNoUnmappedRmTypes()
+    return ret
+
+def TestViolatorAttribute(attribute):
+    global FakeTreble
+    ret = ""
+    if FakeTreble:
+        return ret
+
+    violators = DomainsWithAttribute(attribute)
+    if len(violators) > 0:
+        ret += "SELinux: The following domains violate the Treble ban "
+        ret += "against use of the " + attribute + " attribute: "
+        ret += " ".join(str(x) for x in sorted(violators)) + "\n"
+    return ret
+
+def TestViolatorAttributes():
+    ret = TestViolatorAttribute("binder_in_vendor_violators")
+    ret += TestViolatorAttribute("socket_between_core_and_vendor_violators")
+    ret += TestViolatorAttribute("vendor_executes_system_violators")
+    return ret
+
+###
 # extend OptionParser to allow the same option flag to be used multiple times.
 # This is used to allow multiple file_contexts files and tests to be
 # specified.
@@ -225,18 +303,27 @@
         else:
             Option.take_action(self, action, dest, opt, value, values, parser)
 
-Tests = ["CoredomainViolators"]
+Tests = {"CoredomainViolations": TestCoredomainViolations,
+         "TrebleCompatMapping": TestTrebleCompatMapping,
+         "ViolatorAttributes": TestViolatorAttributes}
 
 if __name__ == '__main__':
-    usage = "sepolicy-trebletests -f nonplat_file_contexts -f "
-    usage +="plat_file_contexts -p policy [--test test] [--help]"
+    usage = "treble_sepolicy_tests.py -l out/host/linux-x86/lib64 "
+    usage += "-f nonplat_file_contexts -f plat_file_contexts "
+    usage += "-p curr_policy -b base_policy -o old_policy "
+    usage +="-m mapping file [--test test] [--help]"
     parser = OptionParser(option_class=MultipleOption, usage=usage)
+    parser.add_option("-b", "--basepolicy", dest="basepolicy", metavar="FILE")
     parser.add_option("-f", "--file_contexts", dest="file_contexts",
             metavar="FILE", action="extend", type="string")
-    parser.add_option("-p", "--policy", dest="policy", metavar="FILE")
     parser.add_option("-l", "--library-path", dest="libpath", metavar="FILE")
-    parser.add_option("-t", "--test", dest="test", action="extend",
+    parser.add_option("-m", "--mapping", dest="mapping", metavar="FILE")
+    parser.add_option("-o", "--oldpolicy", dest="oldpolicy", metavar="FILE")
+    parser.add_option("-p", "--policy", dest="policy", metavar="FILE")
+    parser.add_option("-t", "--test", dest="tests", action="extend",
             help="Test options include "+str(Tests))
+    parser.add_option("--fake-treble", action="store_true", dest="faketreble",
+            default=False)
 
     (options, args) = parser.parse_args()
 
@@ -245,9 +332,14 @@
     if not os.path.exists(options.libpath):
         sys.exit("Error: library-path " + options.libpath + " does not exist\n"
                 + parser.usage)
-
+    if not options.basepolicy:
+        sys.exit("Must specify the current platform-only policy file\n" + parser.usage)
+    if not options.mapping:
+        sys.exit("Must specify a compatibility mapping file\n" + parser.usage)
+    if not options.oldpolicy:
+        sys.exit("Must specify the previous monolithic policy file\n" + parser.usage)
     if not options.policy:
-        sys.exit("Must specify monolithic policy file\n" + parser.usage)
+        sys.exit("Must specify current monolithic policy file\n" + parser.usage)
     if not os.path.exists(options.policy):
         sys.exit("Error: policy file " + options.policy + " does not exist\n"
                 + parser.usage)
@@ -259,16 +351,35 @@
             sys.exit("Error: File_contexts file " + f + " does not exist\n" +
                     parser.usage)
 
+    if options.faketreble:
+        FakeTreble = True
+
     pol = policy.Policy(options.policy, options.file_contexts, options.libpath)
     setup(pol)
+    basepol = policy.Policy(options.basepolicy, None, options.libpath)
+    oldpol = policy.Policy(options.oldpolicy, None, options.libpath)
+    mapping = mini_parser.MiniCilParser(options.mapping)
+    compatSetup(basepol, oldpol, mapping)
 
     if DEBUG:
         PrintScontexts()
 
     results = ""
     # If an individual test is not specified, run all tests.
-    if options.test is None or "CoredomainViolations" in options.tests:
-        results += TestCoredomainViolations()
+    if options.tests is None:
+        for t in Tests.values():
+            results += t()
+    else:
+        for tn in options.tests:
+            t = Tests.get(tn)
+            if t:
+                results += t()
+            else:
+                err = "Error: unknown test: " + tn + "\n"
+                err += "Available tests:\n"
+                for tn in Tests.keys():
+                    err += tn + "\n"
+                sys.exit(err)
 
     if len(results) > 0:
         sys.exit(results)