1# coding: utf-8
2# Copyright 2009 Alexandre Fiori
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import random
17
18from twisted.internet import defer
19from twisted.trial import unittest
20
21import txredisapi as redis
22
23from tests.mixins import REDIS_HOST, REDIS_PORT
24
25
26class SetsTests(unittest.TestCase):
27    '''
28    Tests to ensure that set returning operations return sets
29    '''
30    _KEYS = ['txredisapi:testsets1', 'txredisapi:testsets2',
31             'txredisapi:testsets3', 'txredisapi:testsets4']
32    N = 1024
33
34    @defer.inlineCallbacks
35    def setUp(self):
36        self.db = yield redis.Connection(REDIS_HOST, REDIS_PORT,
37                                         reconnect=False)
38
39    @defer.inlineCallbacks
40    def tearDown(self):
41        yield defer.gatherResults([self.db.delete(x) for x in self._KEYS])
42        yield self.db.disconnect()
43
44    @defer.inlineCallbacks
45    def test_saddrem(self):
46        s = set(range(self.N))
47        r = yield self.db.sadd(self._KEYS[0], s)
48        self.assertEqual(r, len(s))
49        a = s.pop()
50        r = yield self.db.srem(self._KEYS[0], a)
51        self.assertEqual(r, 1)
52        l = [s.pop() for x in range(self.N >> 2)]
53        r = yield self.db.srem(self._KEYS[0], l)
54        self.assertEqual(r, len(l))
55        r = yield self.db.srem(self._KEYS[0], self.N + 1)
56        self.assertEqual(r, 0)
57        r = yield self.db.smembers(self._KEYS[0])
58        self.assertIsInstance(r, set)
59        self.assertEqual(r, s)
60
61    @defer.inlineCallbacks
62    def _test_set(self, key, s):
63        '''
64        Check if the Redis set and the Python set are identical
65        '''
66        r = yield self.db.scard(key)
67        self.assertEqual(r, len(s))
68        r = yield self.db.smembers(key)
69        self.assertEqual(r, s)
70
71    @defer.inlineCallbacks
72    def test_sunion(self):
73        s = set(range(self.N))
74        s1 = set()
75        for x in range(4):
76            ss = set(s.pop() for x in range(self.N >> 2))
77            s1.update(ss)
78            r = yield self.db.sadd(self._KEYS[x], ss)
79            self.assertEqual(r, len(ss))
80        r = yield self.db.sunion(self._KEYS[:4])
81        self.assertIsInstance(r, set)
82        self.assertEqual(r, s1)
83        # Test sunionstore
84        r = yield self.db.sunionstore(self._KEYS[0], self._KEYS[:4])
85        self.assertEqual(r, len(s1))
86        yield self._test_set(self._KEYS[0], s1)
87
88    @defer.inlineCallbacks
89    def test_sdiff(self):
90        l = list(range(self.N))
91        random.shuffle(l)
92        p1 = set(l[:self.N >> 1])
93        random.shuffle(l)
94        p2 = set(l[:self.N >> 1])
95        r = yield self.db.sadd(self._KEYS[0], p1)
96        self.assertEqual(r, len(p1))
97        r = yield self.db.sadd(self._KEYS[1], p2)
98        self.assertEqual(r, len(p2))
99        r = yield self.db.sdiff(self._KEYS[:2])
100        self.assertIsInstance(r, set)
101        a = p1 - p2
102        self.assertEqual(r, a)
103        # Test sdiffstore
104        r = yield self.db.sdiffstore(self._KEYS[0], self._KEYS[:2])
105        self.assertEqual(r, len(a))
106        yield self._test_set(self._KEYS[0], a)
107
108    @defer.inlineCallbacks
109    def test_sinter(self):
110        l = list(range(self.N))
111        random.shuffle(l)
112        p1 = set(l[:self.N >> 1])
113        random.shuffle(l)
114        p2 = set(l[:self.N >> 1])
115        r = yield self.db.sadd(self._KEYS[0], p1)
116        self.assertEqual(r, len(p1))
117        r = yield self.db.sadd(self._KEYS[1], p2)
118        self.assertEqual(r, len(p2))
119        r = yield self.db.sinter(self._KEYS[:2])
120        self.assertIsInstance(r, set)
121        a = p1.intersection(p2)
122        self.assertEqual(r, a)
123        # Test sinterstore
124        r = yield self.db.sinterstore(self._KEYS[0], self._KEYS[:2])
125        self.assertEqual(r, len(a))
126        yield self._test_set(self._KEYS[0], a)
127
128    @defer.inlineCallbacks
129    def test_smembers(self):
130        s = set(range(self.N))
131        r = yield self.db.sadd(self._KEYS[0], s)
132        self.assertEqual(r, len(s))
133        r = yield self.db.smembers(self._KEYS[0])
134        self.assertIsInstance(r, set)
135        self.assertEqual(r, s)
136
137    @defer.inlineCallbacks
138    def test_sismemember(self):
139        yield self.db.sadd(self._KEYS[0], 1)
140        r = yield self.db.sismember(self._KEYS[0], 1)
141        self.assertIsInstance(r, bool)
142        self.assertEqual(r, True)
143        yield self.db.srem(self._KEYS[0], 1)
144        r = yield self.db.sismember(self._KEYS[0], 1)
145        self.assertIsInstance(r, bool)
146        self.assertEqual(r, False)
147
148    @defer.inlineCallbacks
149    def test_smove(self):
150        yield self.db.sadd(self._KEYS[0], [1, 2, 3])
151        # Test moving an existing element
152        r = yield self.db.smove(self._KEYS[0], self._KEYS[1], 1)
153        self.assertIsInstance(r, bool)
154        self.assertEqual(r, True)
155        r = yield self.db.smembers(self._KEYS[1])
156        self.assertEqual(r, set([1]))
157        # Test moving an non existing element
158        r = yield self.db.smove(self._KEYS[0], self._KEYS[1], 4)
159        self.assertIsInstance(r, bool)
160        self.assertEqual(r, False)
161        r = yield self.db.smembers(self._KEYS[1])
162        self.assertEqual(r, set([1]))
163
164    @defer.inlineCallbacks
165    def test_srandmember(self):
166        l = range(10)
167        yield self.db.sadd(self._KEYS[0], l)
168        for i in l:
169            r = yield self.db.srandmember(self._KEYS[0])
170            self.assertIn(r, l)
171
172    @defer.inlineCallbacks
173    def test_spop(self):
174        l = range(10)
175        yield self.db.sadd(self._KEYS[0], l)
176        popped = set()
177        for i in l:
178            r = yield self.db.spop(self._KEYS[0])
179            self.assertNotIn(r, popped)
180            popped.add(r)
181        self.assertEqual(set(l), popped)
182