src/third_party/chromite/lib/paygen/flock_unittest.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Test flock library."""
   7
   8 from __future__ import print_function
   9
  10 import mox
  11 import multiprocessing
  12 import os
  13 import sys
  14 import time
  15
  16 import fixup_path
  17 fixup_path.FixupPath()
  18
  19 from chromite.lib import cros_test_lib
  20 from chromite.lib import osutils
  21 from chromite.lib.paygen import flock
  22
  23
  24 LOCK_ACQUIRED = 5
  25 LOCK_NOT_ACQUIRED = 6
  26
  27
  28 class FLockTest(mox.MoxTestBase):
  29   """Test FLock lock class."""
  30
  31   def __init__(self, testCaseNames):
  32     self.tempdir = None
  33     mox.MoxTestBase.__init__(self, testCaseNames)
  34
  35   def setUp(self):
  36     """Prepare for each test."""
  37     self.mox = mox.Mox()
  38
  39     # To make certain we don't self update while running tests.
  40     os.environ['CROSTOOLS_NO_SOURCE_UPDATE'] = '1'
  41
  42   def tearDown(self):
  43     """Cleanup after each test."""
  44     self.mox.UnsetStubs()
  45
  46   @osutils.TempDirDecorator
  47   def _HelperSingleLockTest(self, blocking, shared):
  48     """Helper method that runs a basic test with/without blocking/sharing."""
  49     lock = flock.Lock('SingleLock',
  50                       lock_dir=self.tempdir,
  51                       blocking=blocking,
  52                       shared=shared)
  53
  54     expected_lock_file = os.path.join(self.tempdir, 'SingleLock')
  55
  56     self.assertFalse(os.path.exists(expected_lock_file))
  57     self.assertFalse(lock.IsLocked())
  58     lock.Acquire()
  59     self.assertTrue(os.path.exists(expected_lock_file))
  60     self.assertTrue(lock.IsLocked())
  61
  62     # Acquiring the lock again should be safe.
  63     lock.Acquire()
  64     self.assertTrue(lock.IsLocked())
  65
  66     # Ensure the lock file contains our pid, and nothing else.
  67     fd = open(expected_lock_file, 'r')
  68     self.assertEquals(['%d\n' % os.getpid()], fd.readlines())
  69     fd.close()
  70
  71     lock.Release()
  72     self.assertFalse(lock.IsLocked())
  73
  74   @osutils.TempDirDecorator
  75   def _HelperDoubleLockTest(self, blocking1, shared1, blocking2, shared2):
  76     """Helper method that runs a two-lock test with/without blocking/sharing."""
  77     lock1 = flock.Lock('DoubleLock',
  78                        lock_dir=self.tempdir,
  79                        blocking=blocking1,
  80                        shared=shared1)
  81     lock2 = flock.Lock('DoubleLock',
  82                        lock_dir=self.tempdir,
  83                        blocking=blocking2,
  84                        shared=shared2)
  85
  86     lock1.Acquire()
  87     self.assertTrue(lock1.IsLocked())
  88     self.assertFalse(lock2.IsLocked())
  89
  90     # The second lock should fail to acquire.
  91     self.assertRaises(flock.LockNotAcquired, lock2.Acquire)
  92     self.assertTrue(lock1.IsLocked())
  93     self.assertFalse(lock2.IsLocked())
  94
  95     lock1.Release()
  96     self.assertFalse(lock1.IsLocked())
  97     self.assertFalse(lock2.IsLocked())
  98
  99     # Releasing second lock should be harmless.
 100     lock2.Release()
 101     self.assertFalse(lock1.IsLocked())
 102     self.assertFalse(lock2.IsLocked())
 103
 104   def _HelperInsideProcess(self, name, lock_dir, blocking, shared):
 105     """Helper method that runs a basic test with/without blocking."""
 106
 107     try:
 108       with flock.Lock(name,
 109                       lock_dir=lock_dir,
 110                       blocking=blocking,
 111                       shared=shared):
 112         pass
 113       sys.exit(LOCK_ACQUIRED)
 114     except flock.LockNotAcquired:
 115       sys.exit(LOCK_NOT_ACQUIRED)
 116
 117   def _HelperStartProcess(self, name, blocking=False, shared=False):
 118     """Create a process and invoke _HelperInsideProcess in it."""
 119     p = multiprocessing.Process(target=self._HelperInsideProcess,
 120                                 args=(name, self.tempdir, blocking, shared))
 121     p.start()
 122
 123     # It's highly probably that p will have tried to grab the lock before the
 124     # timer expired, but not certain.
 125     time.sleep(0.1)
 126
 127     return p
 128
 129   def _HelperWithProcess(self, name, expected, blocking=False, shared=False):
 130     """Create a process and invoke _HelperInsideProcess in it."""
 131     p = multiprocessing.Process(target=self._HelperInsideProcess,
 132                                 args=(name, self.tempdir, blocking, shared))
 133     p.start()
 134     p.join()
 135     self.assertEquals(p.exitcode, expected)
 136
 137   def testLockName(self):
 138     """Make sure that we get the expected lock file name."""
 139     lock = flock.Lock(lock_name='/tmp/foo')
 140     self.assertEqual(lock.lock_file, '/tmp/foo')
 141
 142     lock = flock.Lock(lock_name='foo')
 143     self.assertEqual(lock.lock_file, '/tmp/run_once/foo')
 144
 145     lock = flock.Lock(lock_name='foo', lock_dir='/bar')
 146     self.assertEqual(lock.lock_file, '/bar/foo')
 147
 148   def testSingleLock(self):
 149     """Just test getting releasing a lock with options."""
 150     self._HelperSingleLockTest(blocking=False, shared=False)
 151     self._HelperSingleLockTest(blocking=True, shared=False)
 152     self._HelperSingleLockTest(blocking=False, shared=True)
 153     self._HelperSingleLockTest(blocking=True, shared=True)
 154
 155   def testDoubleLock(self):
 156     """Test two lock objects for the same lock file."""
 157     self._HelperDoubleLockTest(blocking1=False, shared1=False,
 158                                blocking2=False, shared2=False)
 159
 160   def testContextMgr(self):
 161     """Make sure we behave properly with 'with'."""
 162
 163     name = 'WithLock'
 164
 165     # Create an instance, and use it in a with
 166     prelock = flock.Lock(name, lock_dir=self.tempdir)
 167     self._HelperWithProcess(name, expected=LOCK_ACQUIRED)
 168
 169     with prelock as lock:
 170       # Assert the instance didn't change.
 171       self.assertIs(prelock, lock)
 172       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 173
 174     self._HelperWithProcess(name, expected=LOCK_ACQUIRED)
 175
 176     # Construct the instance in the with expression.
 177     with flock.Lock(name, lock_dir=self.tempdir) as lock:
 178       self.assertIsInstance(lock, flock.Lock)
 179       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 180
 181     self._HelperWithProcess(name, expected=LOCK_ACQUIRED)
 182
 183   def testAcquireBeforeWith(self):
 184     """Sometimes you want to Acquire a lock and then return it into 'with'."""
 185
 186     name = "WithLock"
 187     lock = flock.Lock(name, lock_dir=self.tempdir)
 188     lock.Acquire()
 189
 190     self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 191
 192     with lock:
 193       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 194
 195     self._HelperWithProcess(name, expected=LOCK_ACQUIRED)
 196
 197   @osutils.TempDirDecorator
 198   def testSingleProcessLock(self):
 199     """Test grabbing the same lock in processes with no conflicts."""
 200     self._HelperWithProcess('ProcessLock', expected=LOCK_ACQUIRED)
 201     self._HelperWithProcess('ProcessLock', expected=LOCK_ACQUIRED)
 202     self._HelperWithProcess('ProcessLock', expected=LOCK_ACQUIRED,
 203                             blocking=True)
 204     self._HelperWithProcess('ProcessLock', expected=LOCK_ACQUIRED,
 205                             shared=True)
 206     self._HelperWithProcess('ProcessLock', expected=LOCK_ACQUIRED,
 207                             blocking=True, shared=True)
 208
 209   @osutils.TempDirDecorator
 210   def testNonBlockingConflicts(self):
 211     """Test that we get a lock conflict for non-blocking locks."""
 212     name = 'ProcessLock'
 213     with flock.Lock(name, lock_dir=self.tempdir):
 214       self._HelperWithProcess(name,
 215                               expected=LOCK_NOT_ACQUIRED)
 216
 217       self._HelperWithProcess(name,
 218                               expected=LOCK_NOT_ACQUIRED,
 219                               shared=True)
 220
 221     # Can grab it after it's released
 222     self._HelperWithProcess(name, expected=LOCK_ACQUIRED)
 223
 224   @osutils.TempDirDecorator
 225   def testSharedLocks(self):
 226     """Test lock conflict for blocking locks."""
 227     name = 'ProcessLock'
 228
 229     # Intial lock is NOT shared
 230     with flock.Lock(name, lock_dir=self.tempdir, shared=False):
 231       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED, shared=True)
 232
 233     # Intial lock IS shared
 234     with flock.Lock(name, lock_dir=self.tempdir, shared=True):
 235       self._HelperWithProcess(name, expected=LOCK_ACQUIRED, shared=True)
 236       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED, shared=False)
 237
 238   @osutils.TempDirDecorator
 239   def testBlockingConflicts(self):
 240     """Test lock conflict for blocking locks."""
 241     name = 'ProcessLock'
 242
 243     # Intial lock is blocking, exclusive
 244     with flock.Lock(name, lock_dir=self.tempdir, blocking=True):
 245       self._HelperWithProcess(name,
 246                               expected=LOCK_NOT_ACQUIRED,
 247                               blocking=False)
 248
 249       p = self._HelperStartProcess(name, blocking=True, shared=False)
 250
 251     # when the with clause exits, p should unblock and get the lock, setting
 252     # its exit code to sucess now.
 253     p.join()
 254     self.assertEquals(p.exitcode, LOCK_ACQUIRED)
 255
 256     # Intial lock is NON blocking
 257     with flock.Lock(name, lock_dir=self.tempdir):
 258       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 259
 260       p = self._HelperStartProcess(name, blocking=True, shared=False)
 261
 262     # when the with clause exits, p should unblock and get the lock, setting
 263     # it's exit code to sucess now.
 264     p.join()
 265     self.assertEquals(p.exitcode, LOCK_ACQUIRED)
 266
 267     # Intial lock is shared, blocking lock is exclusive
 268     with flock.Lock(name, lock_dir=self.tempdir, shared=True):
 269       self._HelperWithProcess(name, expected=LOCK_NOT_ACQUIRED)
 270       self._HelperWithProcess(name, expected=LOCK_ACQUIRED, shared=True)
 271
 272       p = self._HelperStartProcess(name, blocking=True, shared=False)
 273       q = self._HelperStartProcess(name, blocking=True, shared=False)
 274
 275     # when the with clause exits, p should unblock and get the lock, setting
 276     # it's exit code to sucess now.
 277     p.join()
 278     self.assertEquals(p.exitcode, LOCK_ACQUIRED)
 279     q.join()
 280     self.assertEquals(p.exitcode, LOCK_ACQUIRED)
 281
 282
 283 if __name__ == '__main__':
 284   cros_test_lib.main()