~ubuntu-branches/ubuntu/wily/pymongo/wily

« back to all changes in this revision

Viewing changes to test/test_pooling_gevent.py

  • Committer: Package Import Robot
  • Author(s): Federico Ceratto
  • Date: 2012-05-10 21:21:40 UTC
  • mfrom: (1.1.11)
  • Revision ID: package-import@ubuntu.com-20120510212140-9c66c00zz850h6l9
Tags: 2.2-1
* New upstream release.
* Dependencies added (Closes: #670268)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright 2012 10gen, Inc.
 
2
#
 
3
# Licensed under the Apache License, Version 2.0 (the "License"); you
 
4
# may not use this file except in compliance with the License.  You
 
5
# may obtain a copy of the License at
 
6
#
 
7
# http://www.apache.org/licenses/LICENSE-2.0
 
8
#
 
9
# Unless required by applicable law or agreed to in writing, software
 
10
# distributed under the License is distributed on an "AS IS" BASIS,
 
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 
12
# implied.  See the License for the specific language governing
 
13
# permissions and limitations under the License.
 
14
 
 
15
"""Tests for connection-pooling with greenlets and Gevent"""
 
16
 
 
17
import unittest
 
18
 
 
19
from nose.plugins.skip import SkipTest
 
20
 
 
21
from pymongo import pool
 
22
from test.test_connection import host, port
 
23
from test.test_pooling_base import (
 
24
    _TestPooling, _TestMaxPoolSize, _TestPoolSocketSharing)
 
25
 
 
26
 
 
27
def looplet(greenlets):
 
28
    """World's smallest event loop; run until all greenlets are done
 
29
    """
 
30
    while True:
 
31
        done = True
 
32
 
 
33
        for g in greenlets:
 
34
            if not g.dead:
 
35
                done = False
 
36
                g.switch()
 
37
 
 
38
        if done:
 
39
            return
 
40
 
 
41
 
 
42
class TestPoolingGevent(_TestPooling, unittest.TestCase):
 
43
    """Apply all the standard pool tests to GreenletPool with Gevent"""
 
44
    use_greenlets = True
 
45
 
 
46
 
 
47
class TestPoolingGeventSpecial(unittest.TestCase):
 
48
    """Do a few special GreenletPool tests that don't use TestPoolingBase"""
 
49
    def test_greenlet_sockets(self):
 
50
        # Check that Pool gives two sockets to two greenlets
 
51
        try:
 
52
            import greenlet
 
53
        except ImportError:
 
54
            raise SkipTest('greenlet not installed')
 
55
 
 
56
        cx_pool = pool.GreenletPool(
 
57
            pair=(host,port),
 
58
            max_size=10,
 
59
            net_timeout=1000,
 
60
            conn_timeout=1000,
 
61
            use_ssl=False
 
62
        )
 
63
 
 
64
        socks = []
 
65
 
 
66
        def get_socket():
 
67
            cx_pool.start_request()
 
68
            socks.append(cx_pool.get_socket())
 
69
 
 
70
        looplet([
 
71
            greenlet.greenlet(get_socket),
 
72
            greenlet.greenlet(get_socket),
 
73
        ])
 
74
 
 
75
        self.assertEqual(2, len(socks))
 
76
        self.assertNotEqual(socks[0], socks[1])
 
77
 
 
78
    def test_greenlet_sockets_with_request(self):
 
79
        # Verify two assumptions: that start_request() with two greenlets and
 
80
        # the regular pool will fail, meaning that the two greenlets will
 
81
        # share one socket. Also check that start_request() with GreenletPool
 
82
        # succeeds, meaning that two greenlets will get different sockets (this
 
83
        # is exactly the reason for creating GreenletPool).
 
84
 
 
85
        try:
 
86
            import greenlet
 
87
        except ImportError:
 
88
            raise SkipTest('greenlet not installed')
 
89
 
 
90
        pool_args = dict(
 
91
            pair=(host,port),
 
92
            max_size=10,
 
93
            net_timeout=1000,
 
94
            conn_timeout=1000,
 
95
            use_ssl=False,
 
96
        )
 
97
 
 
98
        for pool_class, use_request, expect_success in [
 
99
            (pool.GreenletPool, True, True),
 
100
            (pool.GreenletPool, False, False),
 
101
            (pool.Pool, True, False),
 
102
            (pool.Pool, False, False),
 
103
        ]:
 
104
            cx_pool = pool_class(**pool_args)
 
105
 
 
106
            # Map: greenlet -> socket
 
107
            greenlet2socks = {}
 
108
            main = greenlet.getcurrent()
 
109
 
 
110
            def get_socket_in_request():
 
111
                # Get a socket from the pool twice, switching contexts each time
 
112
                if use_request:
 
113
                    cx_pool.start_request()
 
114
 
 
115
                main.switch()
 
116
 
 
117
                for _ in range(2):
 
118
                    sock = cx_pool.get_socket()
 
119
                    cx_pool.return_socket(sock)
 
120
                    greenlet2socks.setdefault(
 
121
                        greenlet.getcurrent(), []
 
122
                    ).append(id(sock))
 
123
 
 
124
                    main.switch()
 
125
 
 
126
                cx_pool.end_request()
 
127
 
 
128
            greenlets = [
 
129
                greenlet.greenlet(get_socket_in_request),
 
130
                greenlet.greenlet(get_socket_in_request),
 
131
            ]
 
132
 
 
133
            # Run both greenlets to completion
 
134
            looplet(greenlets)
 
135
 
 
136
            socks_for_gr0 = greenlet2socks[greenlets[0]]
 
137
            socks_for_gr1 = greenlet2socks[greenlets[1]]
 
138
 
 
139
            # Whether we expect requests to work or not, we definitely expect
 
140
            # greenlet2socks to have the same number of keys and values
 
141
            self.assertEqual(2, len(greenlet2socks))
 
142
            self.assertEqual(2, len(socks_for_gr0))
 
143
            self.assertEqual(2, len(socks_for_gr1))
 
144
 
 
145
            # If we started a request, then there was a point at which we had
 
146
            # 2 active sockets, otherwise we always used one.
 
147
            if use_request and pool_class is pool.GreenletPool:
 
148
                self.assertEqual(2, len(cx_pool.sockets))
 
149
            else:
 
150
                self.assertEqual(1, len(cx_pool.sockets))
 
151
 
 
152
            # Again, regardless of whether requests work, a greenlet will get
 
153
            # the same socket each time it calls get_socket() within a request.
 
154
            # What we're really testing is that the two *different* greenlets
 
155
            # get *different* sockets from each other.
 
156
            self.assertEqual(
 
157
                socks_for_gr0[0], socks_for_gr0[1],
 
158
                "Expected greenlet 0 to get the same socket for each call "
 
159
                "to get_socket()"
 
160
            )
 
161
 
 
162
            self.assertEqual(
 
163
                socks_for_gr1[0], socks_for_gr1[1],
 
164
                "Expected greenlet 1 to get the same socket for each call "
 
165
                "to get_socket()"
 
166
            )
 
167
 
 
168
            if expect_success:
 
169
                # We used the proper pool class, so start_request successfully
 
170
                # distinguished between the two greenlets.
 
171
                self.assertNotEqual(
 
172
                    socks_for_gr0[0], socks_for_gr1[0],
 
173
                    "Expected two greenlets to get two different sockets"
 
174
                )
 
175
 
 
176
            else:
 
177
                # We used the wrong pool class, so start_request didn't
 
178
                # distinguish between the two greenlets, and it gave them both
 
179
                # the same socket.
 
180
                self.assertEqual(
 
181
                    socks_for_gr0[0], socks_for_gr1[0],
 
182
                    "Expected two greenlets to get same socket"
 
183
                )
 
184
 
 
185
 
 
186
class TestMaxPoolSizeGevent(_TestMaxPoolSize, unittest.TestCase):
 
187
    use_greenlets = True
 
188
 
 
189
 
 
190
class TestPoolSocketSharingGevent(_TestPoolSocketSharing, unittest.TestCase):
 
191
    use_greenlets = True
 
192
 
 
193
 
 
194
if __name__ == '__main__':
 
195
    unittest.main()