# -*- Python -*-

__version__ = '$Revision: 1.3 $'[11:-2]

from twisted.trial import unittest
from twisted.protocols import htb

class DummyClock:
time = 0
def set(self, when):
self.time = when

def __call__(self):
return self.time

class SomeBucket(htb.Bucket):
maxburst = 100
rate = 2

class TestBucketBase(unittest.TestCase):
def setUp(self):
self._realTimeFunc = htb.time
self.clock = DummyClock()
htb.time = self.clock

def tearDown(self):
htb.time = self._realTimeFunc

class TestBucket(TestBucketBase):
def testBucketSize(self):
"""Testing the size of the bucket."""
b = SomeBucket()
fit = b.add(1000)
self.assertEqual(100, fit)

def testBucketDrain(self):
"""Testing the bucket's drain rate."""
b = SomeBucket()
fit = b.add(1000)
self.clock.set(10)
fit = b.add(1000)
self.assertEqual(20, fit)

def test_bucketEmpty(self):
"""
L{htb.Bucket.drip} returns C{True} if the bucket is empty after that drip.
"""
b = SomeBucket()
b.add(20)
self.clock.set(9)
empty = b.drip()
self.assertFalse(empty)
self.clock.set(10)
empty = b.drip()
self.assertTrue(empty)

class TestBucketNesting(TestBucketBase):
def setUp(self):
TestBucketBase.setUp(self)
self.parent = SomeBucket()
self.child1 = SomeBucket(self.parent)
self.child2 = SomeBucket(self.parent)

def testBucketParentSize(self):
# Use up most of the parent bucket.
self.child1.add(90)
fit = self.child2.add(90)
self.assertEqual(10, fit)

def testBucketParentRate(self):
# Make the parent bucket drain slower.
self.parent.rate = 1
# Fill both child1 and parent.
self.child1.add(100)
self.clock.set(10)
fit = self.child1.add(100)
# How much room was there? The child bucket would have had 20,
# but the parent bucket only ten (so no, it wouldn't make too much
# sense to have a child bucket draining faster than its parent in a real
# application.)
self.assertEqual(10, fit)


# TODO: Test the Transport stuff?

from test_pcp import DummyConsumer

class ConsumerShaperTest(TestBucketBase):
def setUp(self):
TestBucketBase.setUp(self)
self.underlying = DummyConsumer()
self.bucket = SomeBucket()
self.shaped = htb.ShapedConsumer(self.underlying, self.bucket)

def testRate(self):
# Start off with a full bucket, so the burst-size dosen't factor in
# to the calculations.
delta_t = 10
self.bucket.add(100)
self.shaped.write("x" * 100)
self.clock.set(delta_t)
self.shaped.resumeProducing()
self.assertEqual(len(self.underlying.getvalue()),
delta_t * self.bucket.rate)

def testBucketRefs(self):
self.assertEqual(self.bucket._refcount, 1)
self.shaped.stopProducing()
self.assertEqual(self.bucket._refcount, 0)