1 files changed, 20 insertions, 20 deletions
diff --git a/PRESUBMIT_unittest.py b/PRESUBMIT_unittest.py
index 32ff4c3..fef41ed 100755
--- a/PRESUBMIT_unittest.py
+++ b/PRESUBMIT_unittest.py
@@ -14,10 +14,11 @@ import unittest
 
 
 class MockInputApi(object):
-  def __init__(self):
+  def __init__(self, test):
     self.affected_files = []
     self.re = re
     self.os_path = os.path
+    self._test = test
 
   def AffectedFiles(self, include_deletes=True):
     if include_deletes:
@@ -28,6 +29,10 @@ class MockInputApi(object):
   def AffectedTextFiles(self, include_deletes=True):
     return self.affected_files
 
+  def ReadFile(self, file):
+    self._test.failIf(file.LocalPath().endswith('notsource'))
+    return file.file_contents
+
 
 class MockAffectedFile(object):
   def __init__(self, path, action='A'):
@@ -52,19 +57,8 @@ class MockOutputApi(object):
 
 
 class PresubmitUnittest(unittest.TestCase):
-  def setUp(self):
-    self.file_contents = ''
-    def MockReadFile(path):
-      self.failIf(path.endswith('notsource'))
-      return self.file_contents
-    self._ReadFile = PRESUBMIT.ReadFile
-    PRESUBMIT.ReadFile = MockReadFile
-
-  def tearDown(self):
-    PRESUBMIT.ReadFile = self._ReadFile
-
   def testLocalChecks(self):
-    api = MockInputApi()
+    api = MockInputApi(self)
     api.affected_files = [
       MockAffectedFile('foo/blat/yoo.notsource'),
       MockAffectedFile('third_party/blat/source.cc'),
@@ -72,33 +66,39 @@ class PresubmitUnittest(unittest.TestCase):
       MockAffectedFile('foo/blat/source.mm'),
       MockAffectedFile('foo/blat/source.py'),
     ]
-    self.file_contents = 'file with \n\terror\nhere\r\nyes there'
+    for item in api.affected_files:
+      item.file_contents = 'file with \n\terror\nhere\r\nyes there'
     # 3 source files, 2 errors by file + 1 global CR  + 1 global EOF error.
     self.failUnless(len(PRESUBMIT.LocalChecks(api, MockOutputApi)) == 8)
 
-    self.file_contents = 'file\twith\ttabs\n'
+    for item in api.affected_files:
+      item.file_contents = 'file\twith\ttabs\n'
     # 3 source files, 1 error by file.
     self.failUnless(len(PRESUBMIT.LocalChecks(api, MockOutputApi)) == 3)
 
-    self.file_contents = 'file\rusing\rCRs\n'
+    for item in api.affected_files:
+      item.file_contents = 'file\rusing\rCRs\n'
     # One global CR error.
     self.failUnless(len(PRESUBMIT.LocalChecks(api, MockOutputApi)) == 1)
     self.failUnless(
       len(PRESUBMIT.LocalChecks(api, MockOutputApi)[0].items) == 3)
 
-    self.file_contents = 'both\ttabs and\r\nCRLF\n'
+    for item in api.affected_files:
+      item.file_contents = 'both\ttabs and\r\nCRLF\n'
     # 3 source files, 1 error by file + 1 global CR error.
     self.failUnless(len(PRESUBMIT.LocalChecks(api, MockOutputApi)) == 4)
 
-    self.file_contents = 'file with\nzero \\t errors \\r\\n\n'
+    for item in api.affected_files:
+      item.file_contents = 'file with\nzero \\t errors \\r\\n\n'
     self.failIf(PRESUBMIT.LocalChecks(api, MockOutputApi))
 
   def testLocalChecksDeletedFile(self):
-    api = MockInputApi()
+    api = MockInputApi(self)
     api.affected_files = [
       MockAffectedFile('foo/blat/source.py', 'D'),
     ]
-    self.file_contents = 'file with \n\terror\nhere\r\nyes there'
+    api.affected_files[0].file_contents = (
+        'file with \n\terror\nhere\r\nyes there')
     self.failUnless(len(PRESUBMIT.LocalChecks(api, MockOutputApi)) == 0)