apt_dht/PeerManager.py

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