apt_dht/PeerManager.py

   1
   2 from random import choice
   3
   4 from twisted.internet import reactor, defer
   5 from twisted.trial import unittest
   6 from twisted.web2 import stream as stream_mod
   7
   8 from HTTPDownloader import HTTPClientManager
   9
  10 class PeerManager:
  11     def __init__(self):
  12         self.clients = {}
  13
  14     def get(self, location_list, method="GET", modtime=None):
  15         """Download from a list of peers.
  16
  17         @type location_list: C{list} of (C{string}, C{int}, C{string})
  18         @var location_list: a list of the locations where the file can be found
  19         """
  20         host, port, path = choice(location_list)
  21         return self.getPeer(host, port, path, method, modtime)
  22
  23     def getPeer(self, host, port, path, method="GET", modtime=None):
  24         if not port:
  25             port = 80
  26         site = host + ":" + str(port)
  27         if site not in self.clients:
  28             self.clients[site] = HTTPClientManager(host, port)
  29         return self.clients[site].get(path, method, modtime)
  30
  31     def close(self):
  32         for site in self.clients:
  33             self.clients[site].close()
  34         self.clients = {}
  35
  36 class TestPeerManager(unittest.TestCase):
  37     """Unit tests for the PeerManager."""
  38
  39     manager = None
  40     pending_calls = []
  41
  42     def gotResp(self, resp, num, expect):
  43         self.failUnless(resp.code >= 200 and resp.code < 300, "Got a non-200 response: %r" % resp.code)
  44         if expect is not None:
  45             self.failUnless(resp.stream.length == expect, "Length was incorrect, got %r, expected %r" % (resp.stream.length, expect))
  46         def print_(n):
  47             pass
  48         def printdone(n):
  49             pass
  50         stream_mod.readStream(resp.stream, print_).addCallback(printdone)
  51
  52     def test_download(self):
  53         self.manager = PeerManager()
  54         self.timeout = 10
  55
  56         host = 'www.camrdale.org'
  57         d = self.manager.get([(host, 80, '/robots.txt')])
  58         d.addCallback(self.gotResp, 1, 309)
  59         return d
  60
  61     def test_head(self):
  62         self.manager = PeerManager()
  63         self.timeout = 10
  64
  65         host = 'www.camrdale.org'
  66         d = self.manager.get([(host, 80, '/robots.txt')], "HEAD")
  67         d.addCallback(self.gotResp, 1, 0)
  68         return d
  69
  70     def test_multiple_downloads(self):
  71         self.manager = PeerManager()
  72         self.timeout = 120
  73         lastDefer = defer.Deferred()
  74
  75         def newRequest(host, path, num, expect, last=False):
  76             d = self.manager.get([(host, 80, path)])
  77             d.addCallback(self.gotResp, num, expect)
  78             if last:
  79                 d.addBoth(lastDefer.callback)
  80
  81         newRequest('www.camrdale.org', "/", 1, 3433)
  82         newRequest('www.camrdale.org', "/blog/", 2, 37121)
  83         newRequest('www.google.ca', "/", 3, None)
  84         self.pending_calls.append(reactor.callLater(1, newRequest, 'www.sfu.ca', '/', 4, None))
  85         self.pending_calls.append(reactor.callLater(10, newRequest, 'www.camrdale.org', '/wikilink.html', 5, 3084))
  86         self.pending_calls.append(reactor.callLater(30, newRequest, 'www.camrdale.org', '/sitemap.html', 6, 4750))
  87         self.pending_calls.append(reactor.callLater(31, newRequest, 'www.sfu.ca', '/studentcentral/index.html', 7, None))
  88         self.pending_calls.append(reactor.callLater(32, newRequest, 'www.camrdale.org', '/openid.html', 8, 2525))
  89         self.pending_calls.append(reactor.callLater(32, newRequest, 'www.camrdale.org', '/subpage.html', 9, 2381))
  90         self.pending_calls.append(reactor.callLater(62, newRequest, 'www.google.ca', '/intl/en/options/', 0, None, True))
  91         return lastDefer
  92
  93     def tearDown(self):
  94         for p in self.pending_calls:
  95             if p.active():
  96                 p.cancel()
  97         self.pending_calls = []
  98         if self.manager:
  99             self.manager.close()
 100             self.manager = None