diff --git a/adb/adb_utils.cpp b/adb/adb_utils.cpp
index cf5fbc8..d1910f1 100644
--- a/adb/adb_utils.cpp
+++ b/adb/adb_utils.cpp
@@ -320,6 +320,10 @@
 }
 
 std::string GetLogFilePath() {
+    // https://issuetracker.google.com/112588493
+    const char* path = getenv("ANDROID_ADB_LOG_PATH");
+    if (path) return path;
+
 #if defined(_WIN32)
     const char log_name[] = "adb.log";
     WCHAR temp_path[MAX_PATH];
diff --git a/adb/client/commandline.cpp b/adb/client/commandline.cpp
index 11a3dfd..936fd5a 100644
--- a/adb/client/commandline.cpp
+++ b/adb/client/commandline.cpp
@@ -96,9 +96,8 @@
         " version                  show version num\n"
         "\n"
         "networking:\n"
-        " connect HOST[:PORT]      connect to a device via TCP/IP [default port=5555]\n"
-        " disconnect [HOST[:PORT]]\n"
-        "     disconnect from given TCP/IP device [default port=5555], or all\n"
+        " connect HOST[:PORT]      connect to a device via TCP/IP\n"
+        " disconnect [[HOST]:PORT] disconnect from given TCP/IP device, or all\n"
         " forward --list           list all forward socket connections\n"
         " forward [--no-rebind] LOCAL REMOTE\n"
         "     forward socket connection using:\n"
@@ -130,7 +129,7 @@
         "     -a: preserve file timestamp and mode\n"
         " sync [all|data|odm|oem|product|system|system_ext|vendor]\n"
         "     sync a local build from $ANDROID_PRODUCT_OUT to the device (default all)\n"
-        "     -l: list but don't copy\n"
+        "     -l: list files that would be copied, but don't copy them\n"
         "\n"
         "shell:\n"
         " shell [-e ESCAPE] [-n] [-Tt] [-x] [COMMAND...]\n"
@@ -1615,13 +1614,13 @@
         return adb_query_command(query);
     }
     else if (!strcmp(argv[0], "connect")) {
-        if (argc != 2) error_exit("usage: adb connect <host>[:<port>]");
+        if (argc != 2) error_exit("usage: adb connect HOST[:PORT>]");
 
         std::string query = android::base::StringPrintf("host:connect:%s", argv[1]);
         return adb_query_command(query);
     }
     else if (!strcmp(argv[0], "disconnect")) {
-        if (argc > 2) error_exit("usage: adb disconnect [<host>[:<port>]]");
+        if (argc > 2) error_exit("usage: adb disconnect [HOST[:PORT]]");
 
         std::string query = android::base::StringPrintf("host:disconnect:%s",
                                                         (argc == 2) ? argv[1] : "");
diff --git a/adb/daemon/shell_service.cpp b/adb/daemon/shell_service.cpp
index de97068..0fb14c4 100644
--- a/adb/daemon/shell_service.cpp
+++ b/adb/daemon/shell_service.cpp
@@ -222,7 +222,7 @@
 bool Subprocess::ForkAndExec(std::string* error) {
     unique_fd child_stdinout_sfd, child_stderr_sfd;
     unique_fd parent_error_sfd, child_error_sfd;
-    char pts_name[PATH_MAX];
+    const char* pts_name = nullptr;
 
     if (command_.empty()) {
         __android_log_security_bswrite(SEC_TAG_ADB_SHELL_INTERACTIVE, "");
@@ -283,10 +283,22 @@
     cenv.push_back(nullptr);
 
     if (type_ == SubprocessType::kPty) {
-        int fd;
-        pid_ = forkpty(&fd, pts_name, nullptr, nullptr);
+        unique_fd pty_master(posix_openpt(O_RDWR | O_NOCTTY | O_CLOEXEC));
+        if (pty_master == -1) {
+            *error =
+                    android::base::StringPrintf("failed to create pty master: %s", strerror(errno));
+            return false;
+        }
+        if (unlockpt(pty_master.get()) != 0) {
+            *error = android::base::StringPrintf("failed to unlockpt pty master: %s",
+                                                 strerror(errno));
+            return false;
+        }
+
+        pid_ = fork();
+        pts_name = ptsname(pty_master.get());
         if (pid_ > 0) {
-          stdinout_sfd_.reset(fd);
+            stdinout_sfd_ = std::move(pty_master);
         }
     } else {
         if (!CreateSocketpair(&stdinout_sfd_, &child_stdinout_sfd)) {
diff --git a/adb/socket_spec.cpp b/adb/socket_spec.cpp
index 1333724..98468b5 100644
--- a/adb/socket_spec.cpp
+++ b/adb/socket_spec.cpp
@@ -93,7 +93,7 @@
         }
     } else {
         std::string addr(spec.substr(4));
-        port_value = -1;
+        port_value = DEFAULT_ADB_LOCAL_TRANSPORT_PORT;
 
         // FIXME: ParseNetAddress rejects port 0. This currently doesn't hurt, because listening
         //        on an address that isn't 'localhost' is unsupported.
diff --git a/adb/socket_spec_test.cpp b/adb/socket_spec_test.cpp
index f5ec0f1..3a2f60c 100644
--- a/adb/socket_spec_test.cpp
+++ b/adb/socket_spec_test.cpp
@@ -20,38 +20,71 @@
 
 #include <gtest/gtest.h>
 
-TEST(socket_spec, parse_tcp_socket_spec) {
+TEST(socket_spec, parse_tcp_socket_spec_just_port) {
     std::string hostname, error, serial;
     int port;
     EXPECT_TRUE(parse_tcp_socket_spec("tcp:5037", &hostname, &port, &serial, &error));
     EXPECT_EQ("", hostname);
     EXPECT_EQ(5037, port);
     EXPECT_EQ("", serial);
+}
 
-    // Bad ports:
+TEST(socket_spec, parse_tcp_socket_spec_bad_ports) {
+    std::string hostname, error, serial;
+    int port;
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:-1", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:65536", &hostname, &port, &serial, &error));
+}
 
+TEST(socket_spec, parse_tcp_socket_spec_host_and_port) {
+    std::string hostname, error, serial;
+    int port;
     EXPECT_TRUE(parse_tcp_socket_spec("tcp:localhost:1234", &hostname, &port, &serial, &error));
     EXPECT_EQ("localhost", hostname);
     EXPECT_EQ(1234, port);
     EXPECT_EQ("localhost:1234", serial);
+}
 
-    EXPECT_FALSE(parse_tcp_socket_spec("tcp:localhost", &hostname, &port, &serial, &error));
+TEST(socket_spec, parse_tcp_socket_spec_host_no_port) {
+    std::string hostname, error, serial;
+    int port;
+    EXPECT_TRUE(parse_tcp_socket_spec("tcp:localhost", &hostname, &port, &serial, &error));
+    EXPECT_EQ("localhost", hostname);
+    EXPECT_EQ(5555, port);
+    EXPECT_EQ("localhost:5555", serial);
+}
+
+TEST(socket_spec, parse_tcp_socket_spec_host_bad_ports) {
+    std::string hostname, error, serial;
+    int port;
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:localhost:", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:localhost:-1", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:localhost:65536", &hostname, &port, &serial, &error));
+}
 
-    // IPv6:
+TEST(socket_spec, parse_tcp_socket_spec_ipv6_and_port) {
+    std::string hostname, error, serial;
+    int port;
     EXPECT_TRUE(parse_tcp_socket_spec("tcp:[::1]:1234", &hostname, &port, &serial, &error));
     EXPECT_EQ("::1", hostname);
     EXPECT_EQ(1234, port);
     EXPECT_EQ("[::1]:1234", serial);
+}
 
+TEST(socket_spec, parse_tcp_socket_spec_ipv6_no_port) {
+    std::string hostname, error, serial;
+    int port;
+    EXPECT_TRUE(parse_tcp_socket_spec("tcp:::1", &hostname, &port, &serial, &error));
+    EXPECT_EQ("::1", hostname);
+    EXPECT_EQ(5555, port);
+    EXPECT_EQ("[::1]:5555", serial);
+}
+
+TEST(socket_spec, parse_tcp_socket_spec_ipv6_bad_ports) {
+    std::string hostname, error, serial;
+    int port;
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:[::1]", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:[::1]:", &hostname, &port, &serial, &error));
     EXPECT_FALSE(parse_tcp_socket_spec("tcp:[::1]:-1", &hostname, &port, &serial, &error));
-    EXPECT_FALSE(parse_tcp_socket_spec("tcp:::1", &hostname, &port, &serial, &error));
-    EXPECT_FALSE(parse_tcp_socket_spec("tcp:::1:1234", &hostname, &port, &serial, &error));
 }
diff --git a/fs_mgr/fs_mgr_overlayfs.cpp b/fs_mgr/fs_mgr_overlayfs.cpp
index 2f0bf7d..ac15ce4 100644
--- a/fs_mgr/fs_mgr_overlayfs.cpp
+++ b/fs_mgr/fs_mgr_overlayfs.cpp
@@ -1009,7 +1009,7 @@
     Fstab candidates;
     for (const auto& entry : fstab) {
         FstabEntry new_entry = entry;
-        if (!fs_mgr_overlayfs_already_mounted(entry.mount_point) ||
+        if (!fs_mgr_overlayfs_already_mounted(entry.mount_point) &&
             !fs_mgr_wants_overlayfs(&new_entry)) {
             continue;
         }
diff --git a/fs_mgr/tests/adb-remount-test.sh b/fs_mgr/tests/adb-remount-test.sh
index c2a0f33..642f2c1 100755
--- a/fs_mgr/tests/adb-remount-test.sh
+++ b/fs_mgr/tests/adb-remount-test.sh
@@ -15,10 +15,13 @@
 
 adb remount tests
 
---help        This help
---serial      Specify device (must if multiple are present)
---color       Dress output with highlighting colors
---print-time  Report the test duration
+--color                     Dress output with highlighting colors
+--help                      This help
+--no-wait-screen            Do not wait for display screen to settle
+--print-time                Report the test duration
+--serial                    Specify device (must if multiple are present)
+--wait-adb <duration>       adb wait timeout
+--wait-fastboot <duration>  fastboot wait timeout
 
 Conditions:
  - Must be a userdebug build.
@@ -53,6 +56,7 @@
 
 ADB_WAIT=4m
 FASTBOOT_WAIT=2m
+screen_wait=true
 
 ##
 ##  Helper Functions
@@ -185,7 +189,7 @@
 [ "USAGE: adb_cat <file> >stdout
 
 Returns: content of file to stdout with carriage returns skipped,
-         true of the file exists" ]
+         true if the file exists" ]
 adb_cat() {
     local OUTPUT="`adb_sh cat ${1} </dev/null 2>&1`"
     local ret=${?}
@@ -193,6 +197,17 @@
     return ${ret}
 }
 
+[ "USAGE: adb_ls <dirfile> >stdout
+
+Returns: filename or directoru content to stdout with carriage returns skipped,
+         true if the ls had no errors" ]
+adb_ls() {
+    local OUTPUT="`adb_sh ls ${1} </dev/null 2>/dev/null`"
+    local ret=${?}
+    echo "${OUTPUT}" | tr -d '\r'
+    return ${ret}
+}
+
 [ "USAGE: adb_reboot
 
 Returns: true if the reboot command succeeded" ]
@@ -436,6 +451,10 @@
 -n - echo newline at exit
 TIMEOUT - default `format_duration ${wait_for_screen_timeout}`" ]
 wait_for_screen() {
+  if ! ${screen_wait}; then
+    adb_wait
+    return
+  fi
   exit_function=true
   if [ X"-n" = X"${1}" ]; then
     exit_function=echo
@@ -743,6 +762,9 @@
 
 OPTIONS=`getopt --alternative --unquoted \
                 --longoptions help,serial:,colour,color,no-colour,no-color \
+                --longoptions wait-adb:,wait-fastboot: \
+                --longoptions wait-screen,wait-display \
+                --longoptions no-wait-screen,no-wait-display \
                 --longoptions gtest_print_time,print-time \
                 -- "?hs:" ${*}` ||
   ( echo "${USAGE}" >&2 ; false ) ||
@@ -766,9 +788,23 @@
     --no-color | --no-colour)
       color=false
       ;;
+    --no-wait-display | --no-wait-screen)
+      screen_wait=false
+      ;;
+    --wait-display | --wait-screen)
+      screen_wait=true
+      ;;
     --print-time | --gtest_print_time)
       print_time=true
       ;;
+    --wait-adb)
+      ADB_WAIT=${2}
+      shift
+      ;;
+    --wait-fastboot)
+      FASTBOOT_WAIT=${2}
+      shift
+      ;;
     --)
       shift
       break
@@ -1145,10 +1181,14 @@
 
 A="Hello World! $(date)"
 echo "${A}" | adb_sh cat - ">/system/hello"
+echo "${A}" | adb_sh cat - ">/system/priv-app/hello"
 echo "${A}" | adb_sh cat - ">/vendor/hello"
 B="`adb_cat /system/hello`" ||
-  die "sytem hello"
+  die "system hello"
 check_eq "${A}" "${B}" /system before reboot
+B="`adb_cat /system/priv-app/hello`" ||
+  die "system priv-app hello"
+check_eq "${A}" "${B}" /system/priv-app before reboot
 B="`adb_cat /vendor/hello`" ||
   die "vendor hello"
 check_eq "${A}" "${B}" /vendor before reboot
@@ -1230,6 +1270,13 @@
 fi
 B="`adb_cat /system/hello`"
 check_eq "${A}" "${B}" /system after reboot
+# If overlayfs has a nested security problem, this will fail.
+B="`adb_ls /system/`" ||
+  dir "adb ls /system"
+[ X"${B}" != X"${B#*priv-app}" ] ||
+  dir "adb ls /system/priv-app"
+B="`adb_cat /system/priv-app/hello`"
+check_eq "${A}" "${B}" /system/priv-app after reboot
 echo "${GREEN}[       OK ]${NORMAL} /system content remains after reboot" >&2
 # Only root can read vendor if sepolicy permissions are as expected.
 adb_root ||
@@ -1351,6 +1398,12 @@
   fi
   B="`adb_cat /system/hello`"
   check_eq "${A}" "${B}" system after flash vendor
+  B="`adb_ls /system/`" ||
+    dir "adb ls /system"
+  [ X"${B}" != X"${B#*priv-app}" ] ||
+    dir "adb ls /system/priv-app"
+  B="`adb_cat /system/priv-app/hello`"
+  check_eq "${A}" "${B}" system/priv-app after flash vendor
   adb_root ||
     die "adb root"
   B="`adb_cat /vendor/hello`"
@@ -1392,15 +1445,17 @@
 echo "${H}"
 [ ${err} = 0 ] &&
   ( adb_sh rm /vendor/hello </dev/null 2>/dev/null || true ) &&
-  adb_sh rm /system/hello </dev/null ||
+  adb_sh rm /system/hello /system/priv-app/hello </dev/null ||
   ( [ -n "${L}" ] && echo "${L}" && false ) ||
   die -t ${T} "cleanup hello"
 B="`adb_cat /system/hello`"
 check_eq "cat: /system/hello: No such file or directory" "${B}" after rm
+B="`adb_cat /system/priv-app/hello`"
+check_eq "cat: /system/priv-app/hello: No such file or directory" "${B}" after rm
 B="`adb_cat /vendor/hello`"
 check_eq "cat: /vendor/hello: No such file or directory" "${B}" after rm
 
-if [ -n "${scratch_partition}" ]; then
+if ${is_bootloader_fastboot} && [ -n "${scratch_partition}" ]; then
 
   echo "${GREEN}[ RUN      ]${NORMAL} test fastboot flash to ${scratch_partition} recovery" >&2
 
@@ -1542,7 +1597,9 @@
     adb_wait ${ADB_WAIT} ||
     die "adb remount -R"
   if [ "orange" != "`get_property ro.boot.verifiedbootstate`" -o \
-       "2" = "`get_property partition.system.verified`" ]; then
+       "2" = "`get_property partition.system.verified`" ] &&
+     [ -n "`get_property ro.boot.verifiedbootstate`" -o \
+       -n "`get_property partition.system.verified`" ]; then
     die "remount -R command failed to disable verity"
   fi
 
diff --git a/libion/ion.c b/libion/ion.c
index b8de5a4..1ecfc78 100644
--- a/libion/ion.c
+++ b/libion/ion.c
@@ -152,6 +152,8 @@
     ion_user_handle_t handle;
     int ret;
 
+    if (!handle_fd) return -EINVAL;
+
     if (!ion_is_legacy(fd)) {
         struct ion_new_allocation_data data = {
             .len = len,
@@ -201,6 +203,7 @@
     int ret;
     struct ion_heap_query query;
 
+    if (!cnt) return -EINVAL;
     memset(&query, 0, sizeof(query));
 
     ret = ion_ioctl(fd, ION_IOC_HEAP_QUERY, &query);
diff --git a/libion/tests/Android.bp b/libion/tests/Android.bp
index b3fcb3b..d3b4688 100644
--- a/libion/tests/Android.bp
+++ b/libion/tests/Android.bp
@@ -18,18 +18,15 @@
     name: "ion-unit-tests",
     cflags: [
         "-g",
-        "-Wall",
-        "-Werror",
         "-Wno-missing-field-initializers",
     ],
     shared_libs: ["libion"],
     srcs: [
-        "ion_test_fixture.cpp",
         "allocate_test.cpp",
-        "formerly_valid_handle_test.cpp",
-        "invalid_values_test.cpp",
-        "map_test.cpp",
-        "device_test.cpp",
         "exit_test.cpp",
+    	"heap_query.cpp",
+        "invalid_values_test.cpp",
+        "ion_test_fixture.cpp",
+        "map_test.cpp",
     ],
 }
diff --git a/libion/tests/allocate_test.cpp b/libion/tests/allocate_test.cpp
index 3c4524e..5ed01bb 100644
--- a/libion/tests/allocate_test.cpp
+++ b/libion/tests/allocate_test.cpp
@@ -14,95 +14,106 @@
  * limitations under the License.
  */
 
-#include <memory>
 #include <sys/mman.h>
+#include <memory>
 
 #include <gtest/gtest.h>
 
 #include <ion/ion.h>
 #include "ion_test_fixture.h"
 
-class Allocate : public IonAllHeapsTest {
-};
+class Allocate : public IonTest {};
 
-TEST_F(Allocate, Allocate)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Allocate, Allocate) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            ion_user_handle_t handle = 0;
-            ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, 0, &handle));
-            ASSERT_TRUE(handle != 0);
-            ASSERT_EQ(0, ion_free(m_ionFd, handle));
+            int fd;
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &fd));
+            ASSERT_TRUE(fd != 0);
+            ASSERT_EQ(close(fd), 0);  // free the buffer
         }
     }
 }
 
-TEST_F(Allocate, AllocateCached)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Allocate, AllocateCached) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            ion_user_handle_t handle = 0;
-            ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED, &handle));
-            ASSERT_TRUE(handle != 0);
-            ASSERT_EQ(0, ion_free(m_ionFd, handle));
+            int fd;
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), ION_FLAG_CACHED, &fd));
+            ASSERT_TRUE(fd != 0);
+            ASSERT_EQ(close(fd), 0);  // free the buffer
         }
     }
 }
 
-TEST_F(Allocate, AllocateCachedNeedsSync)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Allocate, AllocateCachedNeedsSync) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            ion_user_handle_t handle = 0;
-            ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED_NEEDS_SYNC, &handle));
-            ASSERT_TRUE(handle != 0);
-            ASSERT_EQ(0, ion_free(m_ionFd, handle));
+            int fd;
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id),
+                                      ION_FLAG_CACHED_NEEDS_SYNC, &fd));
+            ASSERT_TRUE(fd != 0);
+            ASSERT_EQ(close(fd), 0);  // free the buffer
         }
     }
 }
 
-TEST_F(Allocate, RepeatedAllocate)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Allocate, RepeatedAllocate) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            ion_user_handle_t handle = 0;
+            int fd;
 
             for (unsigned int i = 0; i < 1024; i++) {
                 SCOPED_TRACE(::testing::Message() << "iteration " << i);
-                ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, 0, &handle));
-                ASSERT_TRUE(handle != 0);
-                ASSERT_EQ(0, ion_free(m_ionFd, handle));
+                ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &fd));
+                ASSERT_TRUE(fd != 0);
+                ASSERT_EQ(close(fd), 0);  // free the buffer
             }
         }
     }
 }
 
-TEST_F(Allocate, Zeroed)
-{
+TEST_F(Allocate, Large) {
+    for (const auto& heap : ion_heaps) {
+        SCOPED_TRACE(::testing::Message()
+                     << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
+        int fd;
+        ASSERT_EQ(-ENOMEM,
+                  ion_alloc_fd(ionfd, 3UL * 1024 * 1024 * 1024, 0, (1 << heap.heap_id), 0, &fd));
+    }
+}
+
+// Make sure all heaps always return zeroed pages
+TEST_F(Allocate, Zeroed) {
     auto zeroes_ptr = std::make_unique<char[]>(4096);
 
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+    for (const auto& heap : ion_heaps) {
+        SCOPED_TRACE(::testing::Message()
+                     << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
         int fds[16];
         for (unsigned int i = 0; i < 16; i++) {
             int map_fd = -1;
 
-            ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, 0, &map_fd));
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, 4096, 0, (1 << heap.heap_id), 0, &map_fd));
             ASSERT_GE(map_fd, 0);
 
-            void *ptr = NULL;
+            void* ptr = NULL;
             ptr = mmap(NULL, 4096, PROT_WRITE, MAP_SHARED, map_fd, 0);
             ASSERT_TRUE(ptr != NULL);
 
@@ -116,13 +127,13 @@
             ASSERT_EQ(0, close(fds[i]));
         }
 
-        int newIonFd = ion_open();
+        int new_ionfd = ion_open();
         int map_fd = -1;
 
-        ASSERT_EQ(0, ion_alloc_fd(newIonFd, 4096, 0, heapMask, 0, &map_fd));
+        ASSERT_EQ(0, ion_alloc_fd(new_ionfd, 4096, 0, (1 << heap.heap_id), 0, &map_fd));
         ASSERT_GE(map_fd, 0);
 
-        void *ptr = NULL;
+        void* ptr = NULL;
         ptr = mmap(NULL, 4096, PROT_READ, MAP_SHARED, map_fd, 0);
         ASSERT_TRUE(ptr != NULL);
 
@@ -130,14 +141,6 @@
 
         ASSERT_EQ(0, munmap(ptr, 4096));
         ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Allocate, Large)
-{
-    for (unsigned int heapMask : m_allHeaps) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        ion_user_handle_t handle = 0;
-        ASSERT_EQ(-ENOMEM, ion_alloc(m_ionFd, 3UL*1024*1024*1024, 0, heapMask, 0, &handle));
+        ASSERT_EQ(0, ion_close(new_ionfd));
     }
 }
diff --git a/libion/tests/device_test.cpp b/libion/tests/device_test.cpp
deleted file mode 100644
index eb3f7b6..0000000
--- a/libion/tests/device_test.cpp
+++ /dev/null
@@ -1,546 +0,0 @@
-/*
- * Copyright (C) 2013 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <fcntl.h>
-#include <memory>
-#include <sys/mman.h>
-#include <sys/stat.h>
-#include <sys/types.h>
-
-#include <linux/ion_test.h>
-
-#include <gtest/gtest.h>
-
-#include <ion/ion.h>
-
-#include "ion_test_fixture.h"
-
-#define ALIGN(x,y) (((x) + ((y) - 1)) & ~((y) - 1))
-
-class Device : public IonAllHeapsTest {
- public:
-    virtual void SetUp();
-    virtual void TearDown();
-    int m_deviceFd;
-    void readDMA(int fd, void *buf, size_t size);
-    void writeDMA(int fd, void *buf, size_t size);
-    void readKernel(int fd, void *buf, size_t size);
-    void writeKernel(int fd, void *buf, size_t size);
-    void blowCache();
-    void dirtyCache(void *ptr, size_t size);
-};
-
-void Device::SetUp()
-{
-    IonAllHeapsTest::SetUp();
-    m_deviceFd = open("/dev/ion-test", O_RDONLY);
-    ASSERT_GE(m_deviceFd, 0);
-}
-
-void Device::TearDown()
-{
-    ASSERT_EQ(0, close(m_deviceFd));
-    IonAllHeapsTest::TearDown();
-}
-
-void Device::readDMA(int fd, void *buf, size_t size)
-{
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, fd));
-    struct ion_test_rw_data ion_test_rw_data = {
-            .ptr = (uint64_t)buf,
-            .offset = 0,
-            .size = size,
-            .write = 0,
-    };
-
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_DMA_MAPPING, &ion_test_rw_data));
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, -1));
-}
-
-void Device::writeDMA(int fd, void *buf, size_t size)
-{
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, fd));
-    struct ion_test_rw_data ion_test_rw_data = {
-            .ptr = (uint64_t)buf,
-            .offset = 0,
-            .size = size,
-            .write = 1,
-    };
-
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_DMA_MAPPING, &ion_test_rw_data));
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, -1));
-}
-
-void Device::readKernel(int fd, void *buf, size_t size)
-{
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, fd));
-    struct ion_test_rw_data ion_test_rw_data = {
-            .ptr = (uint64_t)buf,
-            .offset = 0,
-            .size = size,
-            .write = 0,
-    };
-
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_KERNEL_MAPPING, &ion_test_rw_data));
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, -1));
-}
-
-void Device::writeKernel(int fd, void *buf, size_t size)
-{
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, fd));
-    struct ion_test_rw_data ion_test_rw_data = {
-            .ptr = (uint64_t)buf,
-            .offset = 0,
-            .size = size,
-            .write = 1,
-    };
-
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_KERNEL_MAPPING, &ion_test_rw_data));
-    ASSERT_EQ(0, ioctl(m_deviceFd, ION_IOC_TEST_SET_FD, -1));
-}
-
-void Device::blowCache()
-{
-    const size_t bigger_than_cache = 8*1024*1024;
-    void *buf1 = malloc(bigger_than_cache);
-    void *buf2 = malloc(bigger_than_cache);
-    memset(buf1, 0xaa, bigger_than_cache);
-    memcpy(buf2, buf1, bigger_than_cache);
-    free(buf1);
-    free(buf2);
-}
-
-void Device::dirtyCache(void *ptr, size_t size)
-{
-    /* try to dirty cache lines */
-    for (size_t i = size-1; i > 0; i--) {
-        ((volatile char *)ptr)[i];
-        ((char *)ptr)[i] = i;
-    }
-}
-
-TEST_F(Device, KernelReadCached)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        ((char*)buf)[4096] = 0x12;
-        readKernel(map_fd, buf, 4096);
-        ASSERT_EQ(((char*)buf)[4096], 0x12);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, KernelWriteCached)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeKernel(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMAReadCached)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        readDMA(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMAWriteCached)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeDMA(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, KernelReadCachedNeedsSync)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        ((char*)buf)[4096] = 0x12;
-        readKernel(map_fd, buf, 4096);
-        ASSERT_EQ(((char*)buf)[4096], 0x12);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, KernelWriteCachedNeedsSync)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeKernel(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMAReadCachedNeedsSync)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        ion_sync_fd(m_ionFd, map_fd);
-
-        readDMA(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMAWriteCachedNeedsSync)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeDMA(map_fd, buf, 4096);
-
-        ion_sync_fd(m_ionFd, map_fd);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-TEST_F(Device, KernelRead)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = 0;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        ((char*)buf)[4096] = 0x12;
-        readKernel(map_fd, buf, 4096);
-        ASSERT_EQ(((char*)buf)[4096], 0x12);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, KernelWrite)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = 0;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeKernel(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMARead)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = 0;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        for (int i = 0; i < 4096; i++)
-            ((char *)ptr)[i] = i;
-
-        readDMA(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)buf)[i]);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, DMAWrite)
-{
-    auto alloc_ptr = std::make_unique<char[]>(8192 + 1024);
-    void *buf = (void *)(ALIGN((unsigned long)alloc_ptr.get(), 4096) + 1024);
-
-    for (int i = 0; i < 4096; i++)
-        ((char *)buf)[i] = i;
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = 0;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        writeDMA(map_fd, buf, 4096);
-
-        for (int i = 0; i < 4096; i++)
-            ASSERT_EQ((char)i, ((char *)ptr)[i]) << i;
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
-
-TEST_F(Device, IsCached)
-{
-    auto buf_ptr = std::make_unique<char[]>(4096);
-    void *buf = buf_ptr.get();
-
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        int map_fd = -1;
-        unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
-
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, flags, &map_fd));
-        ASSERT_GE(map_fd, 0);
-
-        void *ptr;
-        ptr = mmap(NULL, 4096, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-        ASSERT_TRUE(ptr != NULL);
-
-        dirtyCache(ptr, 4096);
-
-        readDMA(map_fd, buf, 4096);
-
-        bool same = true;
-        for (int i = 4096-16; i >= 0; i -= 16)
-            if (((char *)buf)[i] != i)
-                same = false;
-        ASSERT_FALSE(same);
-
-        ASSERT_EQ(0, munmap(ptr, 4096));
-        ASSERT_EQ(0, close(map_fd));
-    }
-}
diff --git a/libion/tests/exit_test.cpp b/libion/tests/exit_test.cpp
index cdd3e27..f312389 100644
--- a/libion/tests/exit_test.cpp
+++ b/libion/tests/exit_test.cpp
@@ -22,206 +22,206 @@
 
 #include "ion_test_fixture.h"
 
-class Exit : public IonAllHeapsTest {
-};
+class Exit : public IonTest {};
 
-TEST_F(Exit, WithAlloc)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithAllocFd) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                ion_user_handle_t handle = 0;
+            EXPECT_EXIT(
+                    {
+                        int handle_fd = -1;
 
-                ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, 0, &handle));
-                ASSERT_TRUE(handle != 0);
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
+                        ASSERT_EQ(0,
+                                  ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &handle_fd));
+                        ASSERT_NE(-1, handle_fd);
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
         }
     }
 }
 
-TEST_F(Exit, WithAllocFd)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
-        for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-            SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int handle_fd = -1;
-
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, 0, &handle_fd));
-                ASSERT_NE(-1, handle_fd);
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
-        }
-    }
-}
-
-TEST_F(Exit, WithRepeatedAllocFd)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithRepeatedAllocFd) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
             for (unsigned int i = 0; i < 1024; i++) {
-                SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+                SCOPED_TRACE(::testing::Message()
+                             << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
                 SCOPED_TRACE(::testing::Message() << "size " << size);
-                ASSERT_EXIT({
-                    int handle_fd = -1;
+                ASSERT_EXIT(
+                        {
+                            int handle_fd = -1;
 
-                    ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, 0, &handle_fd));
-                    ASSERT_NE(-1, handle_fd);
-                    exit(0);
-                }, ::testing::ExitedWithCode(0), "")
-                        << "failed on heap " << heapMask
-                        << " and size " << size
-                        << " on iteration " << i;
+                            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0,
+                                                      &handle_fd));
+                            ASSERT_NE(-1, handle_fd);
+                            exit(0);
+                        },
+                        ::testing::ExitedWithCode(0), "")
+                        << "failed on heap " << heap.name << ":" << heap.type << ":" << heap.heap_id
+                        << " and size " << size << " on iteration " << i;
             }
         }
     }
 }
 
-
-TEST_F(Exit, WithMapping)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithMapping) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
 
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, 0, &map_fd));
-                ASSERT_GE(map_fd, 0);
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &map_fd));
+                        ASSERT_GE(map_fd, 0);
 
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
-        }
-    }
-
-}
-
-TEST_F(Exit, WithPartialMapping)
-{
-    static const size_t allocationSizes[] = {64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
-        for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-            SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
-
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, 0, &map_fd));
-                ASSERT_GE(map_fd, 0);
-
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
-
-                ASSERT_EQ(0, munmap(ptr, size / 2));
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
         }
     }
 }
 
-TEST_F(Exit, WithMappingCached)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithPartialMapping) {
+    static const size_t allocationSizes[] = {64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
 
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED, &map_fd));
-                ASSERT_GE(map_fd, 0);
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &map_fd));
+                        ASSERT_GE(map_fd, 0);
 
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
-        }
-    }
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
 
-}
-
-TEST_F(Exit, WithPartialMappingCached)
-{
-    static const size_t allocationSizes[] = {64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
-        for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-            SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
-
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED, &map_fd));
-                ASSERT_GE(map_fd, 0);
-
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
-
-                ASSERT_EQ(0, munmap(ptr, size / 2));
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
+                        ASSERT_EQ(0, munmap(ptr, size / 2));
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
         }
     }
 }
 
-TEST_F(Exit, WithMappingNeedsSync)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithMappingCached) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
 
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC, &map_fd));
-                ASSERT_GE(map_fd, 0);
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id),
+                                                  ION_FLAG_CACHED, &map_fd));
+                        ASSERT_GE(map_fd, 0);
 
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
         }
     }
-
 }
 
-TEST_F(Exit, WithPartialMappingNeedsSync)
-{
-    static const size_t allocationSizes[] = {64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Exit, WithPartialMappingCached) {
+    static const size_t allocationSizes[] = {64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
-            EXPECT_EXIT({
-                int map_fd = -1;
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
 
-                ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC, &map_fd));
-                ASSERT_GE(map_fd, 0);
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id),
+                                                  ION_FLAG_CACHED, &map_fd));
+                        ASSERT_GE(map_fd, 0);
 
-                void *ptr;
-                ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
-                ASSERT_TRUE(ptr != NULL);
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
 
-                ASSERT_EQ(0, munmap(ptr, size / 2));
-                exit(0);
-            }, ::testing::ExitedWithCode(0), "");
+                        ASSERT_EQ(0, munmap(ptr, size / 2));
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
+        }
+    }
+}
+
+TEST_F(Exit, WithMappingNeedsSync) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
+        for (size_t size : allocationSizes) {
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
+            SCOPED_TRACE(::testing::Message() << "size " << size);
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
+
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id),
+                                                  ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC,
+                                                  &map_fd));
+                        ASSERT_GE(map_fd, 0);
+
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
+        }
+    }
+}
+
+TEST_F(Exit, WithPartialMappingNeedsSync) {
+    static const size_t allocationSizes[] = {64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
+        for (size_t size : allocationSizes) {
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
+            SCOPED_TRACE(::testing::Message() << "size " << size);
+            EXPECT_EXIT(
+                    {
+                        int map_fd = -1;
+
+                        ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id),
+                                                  ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC,
+                                                  &map_fd));
+                        ASSERT_GE(map_fd, 0);
+
+                        void* ptr;
+                        ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+                        ASSERT_TRUE(ptr != NULL);
+
+                        ASSERT_EQ(0, munmap(ptr, size / 2));
+                        exit(0);
+                    },
+                    ::testing::ExitedWithCode(0), "");
         }
     }
 }
diff --git a/libion/tests/formerly_valid_handle_test.cpp b/libion/tests/formerly_valid_handle_test.cpp
deleted file mode 100644
index 01ab8f3..0000000
--- a/libion/tests/formerly_valid_handle_test.cpp
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * Copyright (C) 2013 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <sys/mman.h>
-
-#include <gtest/gtest.h>
-
-#include <ion/ion.h>
-
-#include "ion_test_fixture.h"
-
-class FormerlyValidHandle : public IonTest {
- public:
-    virtual void SetUp();
-    virtual void TearDown();
-    ion_user_handle_t m_handle;
-};
-
-void FormerlyValidHandle::SetUp()
-{
-    IonTest::SetUp();
-    ASSERT_EQ(0, ion_alloc(m_ionFd, 4096, 0, 1/* ion_env->m_firstHeap */, 0, &m_handle));
-    ASSERT_TRUE(m_handle != 0);
-    ASSERT_EQ(0, ion_free(m_ionFd, m_handle));
-}
-
-void FormerlyValidHandle::TearDown()
-{
-    m_handle = 0;
-}
-
-TEST_F(FormerlyValidHandle, free)
-{
-	ASSERT_EQ(-EINVAL, ion_free(m_ionFd, m_handle));
-}
-
-TEST_F(FormerlyValidHandle, map)
-{
-    int map_fd;
-    unsigned char *ptr;
-
-    ASSERT_EQ(-EINVAL, ion_map(m_ionFd, m_handle, 4096, PROT_READ, 0, 0, &ptr, &map_fd));
-}
-
-TEST_F(FormerlyValidHandle, share)
-{
-    int share_fd;
-
-    ASSERT_EQ(-EINVAL, ion_share(m_ionFd, m_handle, &share_fd));
-}
diff --git a/libion/tests/heap_query.cpp b/libion/tests/heap_query.cpp
new file mode 100644
index 0000000..bad3bbf
--- /dev/null
+++ b/libion/tests/heap_query.cpp
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+#include "ion_test_fixture.h"
+
+class HeapQuery : public IonTest {};
+
+TEST_F(HeapQuery, AtleastOneHeap) {
+    ASSERT_GT(ion_heaps.size(), 0);
+}
+
+// TODO: Check if we expect some of the default
+// heap types to be present on all devices.
diff --git a/libion/tests/invalid_values_test.cpp b/libion/tests/invalid_values_test.cpp
index 77fea17..48fcd72 100644
--- a/libion/tests/invalid_values_test.cpp
+++ b/libion/tests/invalid_values_test.cpp
@@ -16,171 +16,71 @@
 
 #include <sys/mman.h>
 
+#include <memory>
+#include <vector>
+
 #include <gtest/gtest.h>
 
 #include <ion/ion.h>
-
 #include "ion_test_fixture.h"
 
-class InvalidValues : public IonAllHeapsTest {
- public:
-    virtual void SetUp();
-    virtual void TearDown();
-    ion_user_handle_t m_validHandle;
-    int m_validShareFd;
-    ion_user_handle_t const m_badHandle = -1;
-};
+class InvalidValues : public IonTest {};
 
-void InvalidValues::SetUp()
-{
-    IonAllHeapsTest::SetUp();
-    ASSERT_EQ(0, ion_alloc(m_ionFd, 4096, 0, m_firstHeap, 0, &m_validHandle))
-      << m_ionFd << " " << m_firstHeap;
-    ASSERT_TRUE(m_validHandle != 0);
-    ASSERT_EQ(0, ion_share(m_ionFd, m_validHandle, &m_validShareFd));
-}
-
-void InvalidValues::TearDown()
-{
-    ASSERT_EQ(0, ion_free(m_ionFd, m_validHandle));
-    ASSERT_EQ(0, close(m_validShareFd));
-    m_validHandle = 0;
-    IonAllHeapsTest::TearDown();
-}
-
-TEST_F(InvalidValues, ion_close)
-{
+TEST_F(InvalidValues, ion_close) {
     EXPECT_EQ(-EBADF, ion_close(-1));
 }
 
-TEST_F(InvalidValues, ion_alloc)
-{
-    ion_user_handle_t handle;
-    /* invalid ion_fd */
-    int ret = ion_alloc(0, 4096, 0, m_firstHeap, 0, &handle);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion_fd */
-    EXPECT_EQ(-EBADF, ion_alloc(-1, 4096, 0, m_firstHeap, 0, &handle));
-    /* no heaps */
-    EXPECT_EQ(-ENODEV, ion_alloc(m_ionFd, 4096, 0, 0, 0, &handle));
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        /* zero size */
-        EXPECT_EQ(-EINVAL, ion_alloc(m_ionFd, 0, 0, heapMask, 0, &handle));
-        /* too large size */
-        EXPECT_EQ(-EINVAL, ion_alloc(m_ionFd, -1, 0, heapMask, 0, &handle));
-        /* bad alignment */
-        EXPECT_EQ(-EINVAL, ion_alloc(m_ionFd, 4096, -1, heapMask, 0, &handle));
-        /* NULL handle */
-        EXPECT_EQ(-EINVAL, ion_alloc(m_ionFd, 4096, 0, heapMask, 0, NULL));
-    }
-}
-
-TEST_F(InvalidValues, ion_alloc_fd)
-{
+TEST_F(InvalidValues, ion_alloc_fd) {
     int fd;
-    /* invalid ion_fd */
-    int ret = ion_alloc_fd(0, 4096, 0, m_firstHeap, 0, &fd);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion_fd */
-    EXPECT_EQ(-EBADF, ion_alloc_fd(-1, 4096, 0, m_firstHeap, 0, &fd));
-    /* no heaps */
-    EXPECT_EQ(-ENODEV, ion_alloc_fd(m_ionFd, 4096, 0, 0, 0, &fd));
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-        /* zero size */
-        EXPECT_EQ(-EINVAL, ion_alloc_fd(m_ionFd, 0, 0, heapMask, 0, &fd));
-        /* too large size */
-        EXPECT_EQ(-EINVAL, ion_alloc_fd(m_ionFd, -1, 0, heapMask, 0, &fd));
-        /* bad alignment */
-        EXPECT_EQ(-EINVAL, ion_alloc_fd(m_ionFd, 4096, -1, heapMask, 0, &fd));
-        /* NULL handle */
-        EXPECT_EQ(-EINVAL, ion_alloc_fd(m_ionFd, 4096, 0, heapMask, 0, NULL));
+    // no heaps
+    EXPECT_EQ(-ENODEV, ion_alloc_fd(ionfd, 4096, 0, 0, 0, &fd));
+    for (const auto& heap : ion_heaps) {
+        // invalid ion_fd
+        int ret = ion_alloc_fd(0, 4096, 0, (1 << heap.heap_id), 0, &fd);
+        EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
+        // invalid ion_fd
+        EXPECT_EQ(-EBADF, ion_alloc_fd(-1, 4096, 0, (1 << heap.heap_id), 0, &fd));
+        SCOPED_TRACE(::testing::Message()
+                     << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
+        // zero size
+        EXPECT_EQ(-EINVAL, ion_alloc_fd(ionfd, 0, 0, (1 << heap.heap_id), 0, &fd));
+        // too large size
+        EXPECT_EQ(-EINVAL, ion_alloc_fd(ionfd, -1, 0, (1 << heap.heap_id), 0, &fd));
+        // bad alignment
+        // TODO: Current userspace and kernel code completely ignores alignment. So this
+        // test is going to fail. We need to completely remove alignment from the API.
+        // All memory by default is always page aligned. OR actually pass the alignment
+        // down into the kernel and make kernel respect the alignment.
+        // EXPECT_EQ(-EINVAL, ion_alloc_fd(ionfd, 4096, -1, (1 << heap.heap_id), 0, &fd));
+
+        // NULL fd
+        EXPECT_EQ(-EINVAL, ion_alloc_fd(ionfd, 4096, 0, (1 << heap.heap_id), 0, nullptr));
     }
 }
 
-TEST_F(InvalidValues, ion_free)
-{
-    /* invalid ion fd */
-    int ret = ion_free(0, m_validHandle);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion fd */
-    EXPECT_EQ(-EBADF, ion_free(-1, m_validHandle));
-    /* zero handle */
-    EXPECT_EQ(-EINVAL, ion_free(m_ionFd, 0));
-    /* bad handle */
-    EXPECT_EQ(-EINVAL, ion_free(m_ionFd, m_badHandle));
+TEST_F(InvalidValues, ion_query_heap_cnt) {
+    // NULL count
+    EXPECT_EQ(-EINVAL, ion_query_heap_cnt(ionfd, nullptr));
+
+    int heap_count;
+    // bad fd
+    EXPECT_EQ(-EBADF, ion_query_heap_cnt(-1, &heap_count));
 }
 
-TEST_F(InvalidValues, ion_map)
-{
-    int map_fd;
-    unsigned char *ptr;
+TEST_F(InvalidValues, ion_query_get_heaps) {
+    int heap_count;
+    ASSERT_EQ(0, ion_query_heap_cnt(ionfd, &heap_count));
+    ASSERT_GT(heap_count, 0);
 
-    /* invalid ion fd */
-    int ret = ion_map(0, m_validHandle, 4096, PROT_READ, 0, 0, &ptr, &map_fd);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion fd */
-    EXPECT_EQ(-EBADF, ion_map(-1, m_validHandle, 4096, PROT_READ, 0, 0, &ptr, &map_fd));
-    /* zero handle */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, 0, 4096, PROT_READ, 0, 0, &ptr, &map_fd));
-    /* bad handle */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_badHandle, 4096, PROT_READ, 0, 0, &ptr, &map_fd));
-    /* zero length */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_validHandle, 0, PROT_READ, 0, 0, &ptr, &map_fd));
-    /* bad prot */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_validHandle, 4096, -1, 0, 0, &ptr, &map_fd));
-    /* bad offset */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_validHandle, 4096, PROT_READ, 0, -1, &ptr, &map_fd));
-    /* NULL ptr */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_validHandle, 4096, PROT_READ, 0, 0, NULL, &map_fd));
-    /* NULL map_fd */
-    EXPECT_EQ(-EINVAL, ion_map(m_ionFd, m_validHandle, 4096, PROT_READ, 0, 0, &ptr, NULL));
-}
+    // nullptr buffers, still returns success but without
+    // the ion_heap_data.
+    EXPECT_EQ(0, ion_query_get_heaps(ionfd, heap_count, nullptr));
 
-TEST_F(InvalidValues, ion_share)
-{
-    int share_fd;
+    std::unique_ptr<struct ion_heap_data[]> heaps =
+            std::make_unique<struct ion_heap_data[]>(heap_count);
+    // bad fd
+    EXPECT_EQ(-EBADF, ion_query_get_heaps(-1, heap_count, heaps.get()));
 
-    /* invalid ion fd */
-    int ret = ion_share(0, m_validHandle, &share_fd);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion fd */
-    EXPECT_EQ(-EBADF, ion_share(-1, m_validHandle, &share_fd));
-    /* zero handle */
-    EXPECT_EQ(-EINVAL, ion_share(m_ionFd, 0, &share_fd));
-    /* bad handle */
-    EXPECT_EQ(-EINVAL, ion_share(m_ionFd, m_badHandle, &share_fd));
-    /* NULL share_fd */
-    EXPECT_EQ(-EINVAL, ion_share(m_ionFd, m_validHandle, NULL));
-}
-
-TEST_F(InvalidValues, ion_import)
-{
-    ion_user_handle_t handle;
-
-    /* invalid ion fd */
-    int ret = ion_import(0, m_validShareFd, &handle);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion fd */
-    EXPECT_EQ(-EBADF, ion_import(-1, m_validShareFd, &handle));
-    /* bad share_fd */
-    EXPECT_EQ(-EINVAL, ion_import(m_ionFd, 0, &handle));
-    /* invalid share_fd */
-    EXPECT_EQ(-EBADF, ion_import(m_ionFd, -1, &handle));
-    /* NULL handle */
-    EXPECT_EQ(-EINVAL, ion_import(m_ionFd, m_validShareFd, NULL));
-}
-
-TEST_F(InvalidValues, ion_sync_fd)
-{
-    /* invalid ion fd */
-    int ret = ion_sync_fd(0, m_validShareFd);
-    EXPECT_TRUE(ret == -EINVAL || ret == -ENOTTY);
-    /* invalid ion fd */
-    EXPECT_EQ(-EBADF, ion_sync_fd(-1, m_validShareFd));
-    /* bad share_fd */
-    EXPECT_EQ(-EINVAL, ion_sync_fd(m_ionFd, 0));
-    /* invalid share_fd */
-    EXPECT_EQ(-EBADF, ion_sync_fd(m_ionFd, -1));
+    // invalid heap data pointer
+    EXPECT_EQ(-EFAULT, ion_query_get_heaps(ionfd, heap_count, reinterpret_cast<void*>(0xdeadf00d)));
 }
diff --git a/libion/tests/ion_4.12.h b/libion/tests/ion_4.12.h
new file mode 100644
index 0000000..614510c
--- /dev/null
+++ b/libion/tests/ion_4.12.h
@@ -0,0 +1,50 @@
+/****************************************************************************
+ ****************************************************************************
+ ***
+ ***   This header was automatically generated from a Linux kernel header
+ ***   of the same name, to make information necessary for userspace to
+ ***   call into the kernel available to libc.  It contains only constants,
+ ***   structures, and macros generated from the original header, and thus,
+ ***   contains no copyrightable information.
+ ***
+ ***   To edit the content of this header, modify the corresponding
+ ***   source file (e.g. under external/kernel-headers/original/) then
+ ***   run bionic/libc/kernel/tools/update_all.py
+ ***
+ ***   Any manual change here will be lost the next time this script will
+ ***   be run. You've been warned!
+ ***
+ ****************************************************************************
+ ****************************************************************************/
+#ifndef _UAPI_LINUX_ION_NEW_H
+#define _UAPI_LINUX_ION_NEW_H
+#include <linux/ioctl.h>
+#include <linux/types.h>
+#define ION_NUM_HEAP_IDS (sizeof(unsigned int) * 8)
+struct ion_new_allocation_data {
+  __u64 len;
+  __u32 heap_id_mask;
+  __u32 flags;
+  __u32 fd;
+  __u32 unused;
+};
+#define MAX_HEAP_NAME 32
+struct ion_heap_data {
+  char name[MAX_HEAP_NAME];
+  __u32 type;
+  __u32 heap_id;
+  __u32 reserved0;
+  __u32 reserved1;
+  __u32 reserved2;
+};
+struct ion_heap_query {
+  __u32 cnt;
+  __u32 reserved0;
+  __u64 heaps;
+  __u32 reserved1;
+  __u32 reserved2;
+};
+#define ION_IOC_MAGIC 'I'
+#define ION_IOC_NEW_ALLOC _IOWR(ION_IOC_MAGIC, 0, struct ion_new_allocation_data)
+#define ION_IOC_HEAP_QUERY _IOWR(ION_IOC_MAGIC, 8, struct ion_heap_query)
+#endif
diff --git a/libion/tests/ion_test_fixture.cpp b/libion/tests/ion_test_fixture.cpp
index e20c730..935fe5c 100644
--- a/libion/tests/ion_test_fixture.cpp
+++ b/libion/tests/ion_test_fixture.cpp
@@ -15,59 +15,26 @@
  */
 
 #include <gtest/gtest.h>
-
 #include <ion/ion.h>
 
 #include "ion_test_fixture.h"
 
-IonTest::IonTest() : m_ionFd(-1)
-{
-}
+IonTest::IonTest() : ionfd(-1), ion_heaps() {}
 
 void IonTest::SetUp() {
-    m_ionFd = ion_open();
-    ASSERT_GE(m_ionFd, 0);
+    ionfd = ion_open();
+    ASSERT_GE(ionfd, 0);
+
+    int heap_count;
+    int ret = ion_query_heap_cnt(ionfd, &heap_count);
+    ASSERT_EQ(ret, 0);
+    ASSERT_GT(heap_count, 0);
+
+    ion_heaps.resize(heap_count, {});
+    ret = ion_query_get_heaps(ionfd, heap_count, ion_heaps.data());
+    ASSERT_EQ(ret, 0);
 }
 
 void IonTest::TearDown() {
-    ion_close(m_ionFd);
-}
-
-IonAllHeapsTest::IonAllHeapsTest() :
-        m_firstHeap(0),
-        m_lastHeap(0),
-        m_allHeaps()
-{
-}
-
-void IonAllHeapsTest::SetUp() {
-    int fd = ion_open();
-    ASSERT_GE(fd, 0);
-
-    for (int i = 1; i != 0; i <<= 1) {
-        ion_user_handle_t handle = 0;
-        int ret;
-        ret = ion_alloc(fd, 4096, 0, i, 0, &handle);
-        if (ret == 0 && handle != 0) {
-            ion_free(fd, handle);
-            if (!m_firstHeap) {
-                m_firstHeap = i;
-            }
-            m_lastHeap = i;
-            m_allHeaps.push_back(i);
-        } else {
-            ASSERT_EQ(-ENODEV, ret);
-        }
-    }
-    ion_close(fd);
-
-    EXPECT_NE(0U, m_firstHeap);
-    EXPECT_NE(0U, m_lastHeap);
-
-    RecordProperty("Heaps", m_allHeaps.size());
-    IonTest::SetUp();
-}
-
-void IonAllHeapsTest::TearDown() {
-    IonTest::TearDown();
+    ion_close(ionfd);
 }
diff --git a/libion/tests/ion_test_fixture.h b/libion/tests/ion_test_fixture.h
index 4098214..4f254b8 100644
--- a/libion/tests/ion_test_fixture.h
+++ b/libion/tests/ion_test_fixture.h
@@ -18,29 +18,19 @@
 #define ION_TEST_FIXTURE_H_
 
 #include <gtest/gtest.h>
+#include <vector>
+#include "ion_4.12.h"
 
 using ::testing::Test;
 
 class IonTest : public virtual Test {
- public:
+  public:
     IonTest();
-	virtual ~IonTest() {};
-	virtual void SetUp();
-	virtual void TearDown();
-	int m_ionFd;
-};
-
-class IonAllHeapsTest : public IonTest {
- public:
-    IonAllHeapsTest();
-    virtual ~IonAllHeapsTest() {};
+    virtual ~IonTest(){};
     virtual void SetUp();
     virtual void TearDown();
-
-    unsigned int m_firstHeap;
-    unsigned int m_lastHeap;
-
-    std::vector<unsigned int> m_allHeaps;
+    int ionfd;
+    std::vector<struct ion_heap_data> ion_heaps;
 };
 
 #endif /* ION_TEST_FIXTURE_H_ */
diff --git a/libion/tests/map_test.cpp b/libion/tests/map_test.cpp
index c006dc8..f1b47b7 100644
--- a/libion/tests/map_test.cpp
+++ b/libion/tests/map_test.cpp
@@ -15,61 +15,30 @@
  */
 
 #include <sys/mman.h>
+#include <unistd.h>
 
 #include <gtest/gtest.h>
 
 #include <ion/ion.h>
-
 #include "ion_test_fixture.h"
 
-class Map : public IonAllHeapsTest {
-};
+class Map : public IonTest {};
 
-TEST_F(Map, MapHandle)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Map, MapFd) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
-            SCOPED_TRACE(::testing::Message() << "size " << size);
-            ion_user_handle_t handle = 0;
-
-            ASSERT_EQ(0, ion_alloc(m_ionFd, size, 0, heapMask, 0, &handle));
-            ASSERT_TRUE(handle != 0);
-
-            int map_fd = -1;
-            unsigned char *ptr = NULL;
-            ASSERT_EQ(0, ion_map(m_ionFd, handle, size, PROT_READ | PROT_WRITE, MAP_SHARED, 0, &ptr, &map_fd));
-            ASSERT_TRUE(ptr != NULL);
-            ASSERT_GE(map_fd, 0);
-
-            ASSERT_EQ(0, close(map_fd));
-
-            ASSERT_EQ(0, ion_free(m_ionFd, handle));
-
-            memset(ptr, 0xaa, size);
-
-            ASSERT_EQ(0, munmap(ptr, size));
-        }
-    }
-}
-
-TEST_F(Map, MapFd)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
-        for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
             int map_fd = -1;
 
-            ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, 0, &map_fd));
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), 0, &map_fd));
             ASSERT_GE(map_fd, 0);
 
-            void *ptr;
+            void* ptr;
             ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
             ASSERT_TRUE(ptr != NULL);
-
             ASSERT_EQ(0, close(map_fd));
 
             memset(ptr, 0xaa, size);
@@ -79,53 +48,51 @@
     }
 }
 
-TEST_F(Map, MapOffset)
-{
-    for (unsigned int heapMask : m_allHeaps) {
-        SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+TEST_F(Map, MapOffset) {
+    for (const auto& heap : ion_heaps) {
+        SCOPED_TRACE(::testing::Message()
+                     << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
         int map_fd = -1;
 
-        ASSERT_EQ(0, ion_alloc_fd(m_ionFd, PAGE_SIZE * 2, 0, heapMask, 0, &map_fd));
+        ASSERT_EQ(0, ion_alloc_fd(ionfd, getpagesize() * 2, 0, (1 << heap.heap_id), 0, &map_fd));
         ASSERT_GE(map_fd, 0);
 
-        unsigned char *ptr;
-        ptr = (unsigned char *)mmap(NULL, PAGE_SIZE * 2, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
+        unsigned char* ptr;
+        ptr = (unsigned char*)mmap(NULL, getpagesize() * 2, PROT_READ | PROT_WRITE, MAP_SHARED,
+                                   map_fd, 0);
         ASSERT_TRUE(ptr != NULL);
 
-        memset(ptr, 0, PAGE_SIZE);
-        memset(ptr + PAGE_SIZE, 0xaa, PAGE_SIZE);
+        memset(ptr, 0, getpagesize());
+        memset(ptr + getpagesize(), 0xaa, getpagesize());
 
-        ASSERT_EQ(0, munmap(ptr, PAGE_SIZE * 2));
+        ASSERT_EQ(0, munmap(ptr, getpagesize() * 2));
 
-        ptr = (unsigned char *)mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, PAGE_SIZE);
+        ptr = (unsigned char*)mmap(NULL, getpagesize(), PROT_READ | PROT_WRITE, MAP_SHARED, map_fd,
+                                   getpagesize());
         ASSERT_TRUE(ptr != NULL);
-
         ASSERT_EQ(ptr[0], 0xaa);
-        ASSERT_EQ(ptr[PAGE_SIZE - 1], 0xaa);
-
-        ASSERT_EQ(0, munmap(ptr, PAGE_SIZE));
-
+        ASSERT_EQ(ptr[getpagesize() - 1], 0xaa);
+        ASSERT_EQ(0, munmap(ptr, getpagesize()));
         ASSERT_EQ(0, close(map_fd));
     }
 }
 
-TEST_F(Map, MapCached)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Map, MapCached) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
             int map_fd = -1;
             unsigned int flags = ION_FLAG_CACHED;
 
-            ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, flags, &map_fd));
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), flags, &map_fd));
             ASSERT_GE(map_fd, 0);
 
-            void *ptr;
+            void* ptr;
             ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
             ASSERT_TRUE(ptr != NULL);
-
             ASSERT_EQ(0, close(map_fd));
 
             memset(ptr, 0xaa, size);
@@ -135,23 +102,22 @@
     }
 }
 
-TEST_F(Map, MapCachedNeedsSync)
-{
-    static const size_t allocationSizes[] = {4*1024, 64*1024, 1024*1024, 2*1024*1024};
-    for (unsigned int heapMask : m_allHeaps) {
+TEST_F(Map, MapCachedNeedsSync) {
+    static const size_t allocationSizes[] = {4 * 1024, 64 * 1024, 1024 * 1024, 2 * 1024 * 1024};
+    for (const auto& heap : ion_heaps) {
         for (size_t size : allocationSizes) {
-            SCOPED_TRACE(::testing::Message() << "heap " << heapMask);
+            SCOPED_TRACE(::testing::Message()
+                         << "heap:" << heap.name << ":" << heap.type << ":" << heap.heap_id);
             SCOPED_TRACE(::testing::Message() << "size " << size);
             int map_fd = -1;
             unsigned int flags = ION_FLAG_CACHED | ION_FLAG_CACHED_NEEDS_SYNC;
 
-            ASSERT_EQ(0, ion_alloc_fd(m_ionFd, size, 0, heapMask, flags, &map_fd));
+            ASSERT_EQ(0, ion_alloc_fd(ionfd, size, 0, (1 << heap.heap_id), flags, &map_fd));
             ASSERT_GE(map_fd, 0);
 
-            void *ptr;
+            void* ptr;
             ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, map_fd, 0);
             ASSERT_TRUE(ptr != NULL);
-
             ASSERT_EQ(0, close(map_fd));
 
             memset(ptr, 0xaa, size);