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