diff --git a/search/test/ast2ast_test.py b/search/test/ast2ast_test.py
index 9edeaf1..cd50b18 100644
--- a/search/test/ast2ast_test.py
+++ b/search/test/ast2ast_test.py
@@ -1,17 +1,17 @@
-# Copyright 2016 The Chromium Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style
-# license that can be found in the LICENSE file or at
-# https://developers.google.com/open-source/licenses/bsd
+# Copyright 2016 The Chromium Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
 
 """Tests for the ast2ast module."""
 from __future__ import print_function
 from __future__ import division
 from __future__ import absolute_import
 
+import six
 import unittest
 
-from proto import ast_pb2
-from proto import tracker_pb2
+from mrproto import ast_pb2
+from mrproto import tracker_pb2
 from search import ast2ast
 from search import query2ast
 from services import service_manager
@@ -182,7 +182,7 @@
       self.assertEqual(
           'Searching for issues accross multiple/all projects without '
           'project prefixes is ambiguous and is currently not supported.',
-          cm.exception.message)
+          str(cm.exception))
 
   def testPreprocessBlockedOnCond_WithExternalIssues(self):
     blockedon_field = BUILTIN_ISSUE_FIELDS['blockedon']
@@ -322,7 +322,7 @@
       self.assertEqual(
           'Searching for issues accross multiple/all projects without '
           'project prefixes is ambiguous and is currently not supported.',
-          cm.exception.message)
+          str(cm.exception))
 
   def testPreprocessBlockingCond_WithExternalIssues(self):
     blocking_field = BUILTIN_ISSUE_FIELDS['blocking']
@@ -452,19 +452,19 @@
         ast_pb2.QueryOp.IS_DEFINED, [BUILTIN_ISSUE_FIELDS['label']],
         ['Priority', 'Severity'], [])
     regex = ast2ast._MakePrefixRegex(cond)
-    self.assertRegexpMatches('Priority-1', regex)
-    self.assertRegexpMatches('Severity-3', regex)
-    self.assertNotRegexpMatches('My-Priority', regex)
+    self.assertRegex('Priority-1', regex)
+    self.assertRegex('Severity-3', regex)
+    self.assertNotRegex('My-Priority', regex)
 
   def testKeyValueRegex(self):
     cond = ast_pb2.MakeCond(
         ast_pb2.QueryOp.KEY_HAS, [BUILTIN_ISSUE_FIELDS['label']],
         ['Type-Feature', 'Type-Security'], [])
     regex = ast2ast._MakeKeyValueRegex(cond)
-    self.assertRegexpMatches('Type-Feature', regex)
-    self.assertRegexpMatches('Type-Bug-Security', regex)
-    self.assertNotRegexpMatches('Type-Bug', regex)
-    self.assertNotRegexpMatches('Security-Feature', regex)
+    self.assertRegex('Type-Feature', regex)
+    self.assertRegex('Type-Bug-Security', regex)
+    self.assertNotRegex('Type-Bug', regex)
+    self.assertNotRegex('Security-Feature', regex)
 
   def testKeyValueRegex_multipleKeys(self):
     cond = ast_pb2.MakeCond(
@@ -478,8 +478,8 @@
         ast_pb2.QueryOp.TEXT_HAS, [BUILTIN_ISSUE_FIELDS['label']],
         ['Type-Bug'], [])
     regex = ast2ast._MakeKeyValueRegex(cond)
-    self.assertRegexpMatches('Type-Bug-Security', regex)
-    self.assertNotRegexpMatches('Type-BugSecurity', regex)
+    self.assertRegex('Type-Bug-Security', regex)
+    self.assertNotRegex('Type-BugSecurity', regex)
 
   def testPreprocessLabelCond(self):
     label_field = BUILTIN_ISSUE_FIELDS['label']
@@ -693,7 +693,7 @@
         self.cnxn, cond, [1], self.services, None, True)
     self.assertEqual(ast_pb2.QueryOp.EQ, actual.op)
     self.assertEqual([hotlist_id_field], actual.field_defs)
-    self.assertItemsEqual([10, 30, 40, 50, 60], actual.int_values)
+    six.assertCountEqual(self, [10, 30, 40, 50, 60], actual.int_values)
 
   def testPreprocessHotlistCond_UserNotFound(self):
     hotlist_field = BUILTIN_ISSUE_FIELDS['hotlist']