Statistics
| Branch: | Tag: | Revision:

root / test / ganeti.workerpool_unittest.py @ d4fa5c23

History | View | Annotate | Download (3.4 kB)

1 76094e37 Michael Hanselmann
#!/usr/bin/python
2 76094e37 Michael Hanselmann
#
3 76094e37 Michael Hanselmann
4 76094e37 Michael Hanselmann
# Copyright (C) 2008 Google Inc.
5 76094e37 Michael Hanselmann
#
6 76094e37 Michael Hanselmann
# This program is free software; you can redistribute it and/or modify
7 76094e37 Michael Hanselmann
# it under the terms of the GNU General Public License as published by
8 76094e37 Michael Hanselmann
# the Free Software Foundation; either version 2 of the License, or
9 76094e37 Michael Hanselmann
# (at your option) any later version.
10 76094e37 Michael Hanselmann
#
11 76094e37 Michael Hanselmann
# This program is distributed in the hope that it will be useful, but
12 76094e37 Michael Hanselmann
# WITHOUT ANY WARRANTY; without even the implied warranty of
13 76094e37 Michael Hanselmann
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14 76094e37 Michael Hanselmann
# General Public License for more details.
15 76094e37 Michael Hanselmann
#
16 76094e37 Michael Hanselmann
# You should have received a copy of the GNU General Public License
17 76094e37 Michael Hanselmann
# along with this program; if not, write to the Free Software
18 76094e37 Michael Hanselmann
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
19 76094e37 Michael Hanselmann
# 02110-1301, USA.
20 76094e37 Michael Hanselmann
21 76094e37 Michael Hanselmann
22 76094e37 Michael Hanselmann
"""Script for unittesting the workerpool module"""
23 76094e37 Michael Hanselmann
24 76094e37 Michael Hanselmann
import unittest
25 76094e37 Michael Hanselmann
import threading
26 76094e37 Michael Hanselmann
import time
27 76094e37 Michael Hanselmann
import sys
28 76094e37 Michael Hanselmann
import zlib
29 76094e37 Michael Hanselmann
30 76094e37 Michael Hanselmann
from ganeti import workerpool
31 76094e37 Michael Hanselmann
32 76094e37 Michael Hanselmann
33 76094e37 Michael Hanselmann
class DummyBaseWorker(workerpool.BaseWorker):
34 76094e37 Michael Hanselmann
  def RunTask(self, text):
35 76094e37 Michael Hanselmann
    pass
36 76094e37 Michael Hanselmann
37 76094e37 Michael Hanselmann
38 76094e37 Michael Hanselmann
class ChecksumContext:
39 76094e37 Michael Hanselmann
  CHECKSUM_START = zlib.adler32("")
40 76094e37 Michael Hanselmann
41 76094e37 Michael Hanselmann
  def __init__(self):
42 76094e37 Michael Hanselmann
    self.lock = threading.Condition(threading.Lock())
43 76094e37 Michael Hanselmann
    self.checksum = self.CHECKSUM_START
44 76094e37 Michael Hanselmann
45 76094e37 Michael Hanselmann
  @staticmethod
46 76094e37 Michael Hanselmann
  def UpdateChecksum(current, value):
47 76094e37 Michael Hanselmann
    return zlib.adler32(str(value), current)
48 76094e37 Michael Hanselmann
49 76094e37 Michael Hanselmann
50 76094e37 Michael Hanselmann
class ChecksumBaseWorker(workerpool.BaseWorker):
51 76094e37 Michael Hanselmann
  def RunTask(self, ctx, number):
52 76094e37 Michael Hanselmann
    ctx.lock.acquire()
53 76094e37 Michael Hanselmann
    try:
54 76094e37 Michael Hanselmann
      ctx.checksum = ctx.UpdateChecksum(ctx.checksum, number)
55 76094e37 Michael Hanselmann
    finally:
56 76094e37 Michael Hanselmann
      ctx.lock.release()
57 76094e37 Michael Hanselmann
58 76094e37 Michael Hanselmann
59 76094e37 Michael Hanselmann
class TestWorkerpool(unittest.TestCase):
60 76094e37 Michael Hanselmann
  """Workerpool tests"""
61 76094e37 Michael Hanselmann
62 76094e37 Michael Hanselmann
  def testDummy(self):
63 76094e37 Michael Hanselmann
    wp = workerpool.WorkerPool(3, DummyBaseWorker)
64 76094e37 Michael Hanselmann
    try:
65 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 3)
66 76094e37 Michael Hanselmann
67 76094e37 Michael Hanselmann
      for i in xrange(10):
68 76094e37 Michael Hanselmann
        wp.AddTask("Hello world %s" % i)
69 76094e37 Michael Hanselmann
70 76094e37 Michael Hanselmann
      wp.Quiesce()
71 76094e37 Michael Hanselmann
    finally:
72 76094e37 Michael Hanselmann
      wp.TerminateWorkers()
73 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 0)
74 76094e37 Michael Hanselmann
75 76094e37 Michael Hanselmann
  def testNoTasks(self):
76 76094e37 Michael Hanselmann
    wp = workerpool.WorkerPool(3, DummyBaseWorker)
77 76094e37 Michael Hanselmann
    try:
78 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 3)
79 76094e37 Michael Hanselmann
      self._CheckNoTasks(wp)
80 76094e37 Michael Hanselmann
    finally:
81 76094e37 Michael Hanselmann
      wp.TerminateWorkers()
82 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 0)
83 76094e37 Michael Hanselmann
84 76094e37 Michael Hanselmann
  def testNoTasksQuiesce(self):
85 76094e37 Michael Hanselmann
    wp = workerpool.WorkerPool(3, DummyBaseWorker)
86 76094e37 Michael Hanselmann
    try:
87 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 3)
88 76094e37 Michael Hanselmann
      self._CheckNoTasks(wp)
89 76094e37 Michael Hanselmann
      wp.Quiesce()
90 76094e37 Michael Hanselmann
      self._CheckNoTasks(wp)
91 76094e37 Michael Hanselmann
    finally:
92 76094e37 Michael Hanselmann
      wp.TerminateWorkers()
93 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 0)
94 76094e37 Michael Hanselmann
95 76094e37 Michael Hanselmann
  def testChecksum(self):
96 76094e37 Michael Hanselmann
    # Tests whether all tasks are run and, since we're only using a single
97 76094e37 Michael Hanselmann
    # thread, whether everything is started in order.
98 76094e37 Michael Hanselmann
    wp = workerpool.WorkerPool(1, ChecksumBaseWorker)
99 76094e37 Michael Hanselmann
    try:
100 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 1)
101 76094e37 Michael Hanselmann
102 76094e37 Michael Hanselmann
      ctx = ChecksumContext()
103 76094e37 Michael Hanselmann
      checksum = ChecksumContext.CHECKSUM_START
104 76094e37 Michael Hanselmann
      for i in xrange(1, 100):
105 76094e37 Michael Hanselmann
        checksum = ChecksumContext.UpdateChecksum(checksum, i)
106 76094e37 Michael Hanselmann
        wp.AddTask(ctx, i)
107 76094e37 Michael Hanselmann
108 76094e37 Michael Hanselmann
      wp.Quiesce()
109 76094e37 Michael Hanselmann
110 76094e37 Michael Hanselmann
      self._CheckNoTasks(wp)
111 76094e37 Michael Hanselmann
112 76094e37 Michael Hanselmann
      # Check sum
113 76094e37 Michael Hanselmann
      ctx.lock.acquire()
114 76094e37 Michael Hanselmann
      try:
115 76094e37 Michael Hanselmann
        self.assertEqual(checksum, ctx.checksum)
116 76094e37 Michael Hanselmann
      finally:
117 76094e37 Michael Hanselmann
        ctx.lock.release()
118 76094e37 Michael Hanselmann
    finally:
119 76094e37 Michael Hanselmann
      wp.TerminateWorkers()
120 76094e37 Michael Hanselmann
      self._CheckWorkerCount(wp, 0)
121 76094e37 Michael Hanselmann
122 76094e37 Michael Hanselmann
  def _CheckNoTasks(self, wp):
123 76094e37 Michael Hanselmann
    wp._lock.acquire()
124 76094e37 Michael Hanselmann
    try:
125 76094e37 Michael Hanselmann
      # The task queue must be empty now
126 76094e37 Michael Hanselmann
      self.failUnless(not wp._tasks)
127 76094e37 Michael Hanselmann
    finally:
128 76094e37 Michael Hanselmann
      wp._lock.release()
129 76094e37 Michael Hanselmann
130 76094e37 Michael Hanselmann
  def _CheckWorkerCount(self, wp, num_workers):
131 76094e37 Michael Hanselmann
    wp._lock.acquire()
132 76094e37 Michael Hanselmann
    try:
133 76094e37 Michael Hanselmann
      self.assertEqual(len(wp._workers), num_workers)
134 76094e37 Michael Hanselmann
    finally:
135 76094e37 Michael Hanselmann
      wp._lock.release()
136 76094e37 Michael Hanselmann
137 76094e37 Michael Hanselmann
138 76094e37 Michael Hanselmann
if __name__ == '__main__':
139 76094e37 Michael Hanselmann
  unittest.main()