def setUp(self):
Segment.helper = SegmentHelper(redis_obj=fakeredis.FakeStrictRedis(
charset='utf-8', decode_responses=True))
self.u = UserFactory()
self.s = AllUserSegmentFactory()
app_settings.SEGMENTS_REFRESH_ASYNC = False
app_settings.SEGMENTS_REFRESH_ON_SAVE = True
def setUp(self):
helper = SegmentHelper(redis_obj=fakeredis.FakeStrictRedis(
charset='utf-8', decode_responses=True))
Segment.helper = helper
helper.redis.flushdb()
def delete_segment(segment_id):
"""Celery task to delete an individual Segment from Redis """
SegmentHelper().delete_segment(segment_id)
def setUp(self):
self.helper = SegmentHelper()
self.user = UserFactory()
SegmentHelper.redis = fakeredis.FakeStrictRedis(charset='utf-8',
decode_responses=True)
class TestSegmentHelper(TestCase):
def setUp(self):
self.helper = SegmentHelper()
self.user = UserFactory()
SegmentHelper.redis = fakeredis.FakeStrictRedis(charset='utf-8',
decode_responses=True)
def test_add_segment_membership(self):
s = SegmentFactory()
s.add_member(self.user)
self.assertTrue(self.helper.segment_has_member(s.id, self.user.id))
def test_segment_has_member_when_segment_exists(self):
s = SegmentFactory()
s.add_member(self.user)
self.assertTrue(self.helper.segment_has_member(s.id, self.user.id))
def test_segment_has_member_nonexistant_segment(self):
s = SegmentFactory()
s.add_member(self.user)
self.helper.remove_segment_membership(99999, self.user.id)
self.assertTrue(self.helper.segment_has_member(s.id, self.user.id))
def test_remove_segment_membership_segment_exists(self):
s = SegmentFactory()
s.add_member(self.user)
self.helper.remove_segment_membership(s.id, self.user.id)
self.assertFalse(self.helper.segment_has_member(s.id, self.user.id))
def test_get_user_segments_when_segment_exists(self):
s = SegmentFactory()
s.add_member(self.user)
segments = self.helper.get_user_segments(self.user.id)
self.assertTrue(len(segments) > 0)
def test_get_user_segments_when_invalid_user(self):
s = SegmentFactory()
s.add_member(self.user)
segments = self.helper.get_user_segments(9999)
self.assertEquals(len(segments), 0)
def test_get_segment_membership_count(self):
s = SegmentFactory()
s.add_member(self.user)
self.assertEquals(self.helper.get_segment_membership_count(s.id), 1)
def test_get_segment_members_valid_segment(self):
s = SegmentFactory()
s.add_member(self.user)
members = self.helper.get_segment_members(s.id)
self.assertEquals(len(list(members)), 1)
def test_get_segment_members_invalid_segment(self):
s = SegmentFactory()
s.add_member(self.user)
members = self.helper.get_segment_members(99999)
self.assertEquals(len(list(members)), 0)
def test_get_refreshed_users(self):
s = SegmentFactory()
self.helper.refresh_segment(
s.id, 'select %s from %s' % (self.user.pk, user_table()))
self.assertEquals(len(list(self.helper.get_refreshed_users())), 1)
def test_remove_refreshed_user(self):
s = SegmentFactory()
self.helper.refresh_segment(
s.id, 'select %s from %s' % (self.user.pk, user_table()))
self.helper.remove_refreshed_user(self.user.id)
self.assertEquals(len(list(self.helper.get_refreshed_users())), 0)
def test_refresh_segment_invalid_sql(self):
s = SegmentFactory()
invalid_sql = 'abc select '
self.assertEquals(self.helper.refresh_segment(s.id, invalid_sql), 0)
def test_refresh_segment_valid_sql(self):
s = SegmentFactory()
valid_sql = 'select * from %s' % user_table()
self.assertEquals(self.helper.refresh_segment(s.id, valid_sql), 1)
@patch('segments.tasks.delete_segment.delay')
def test_delete_segment(self, p_delete_segment):
s = SegmentFactory()
s.add_member(self.user)
s.delete()
self.assertFalse(self.helper.segment_has_member(s.id, self.user.id))
self.assertTrue(p_delete_segment.called)
def test_diff_segment(self):
s1 = SegmentFactory()
u1 = UserFactory()
s1.add_member(u1)
s2 = SegmentFactory()
u2 = UserFactory()
s2.add_member(u1)
s2.add_member(u2)
s1_key = self.helper.segment_key % s1.id
s2_key = self.helper.segment_key % s2.id
self.helper.diff_segment(s2_key, s1_key, 'diff_test')
self.assertEquals(self.helper.redis.smembers('diff_test'),
{str(u2.id)})
def test_chunk_items(self):
members = [1, 2, 3]
for i in members:
self.assertEquals(
len(list(chunk_items(members, len(members), i))[0]), i)
self.assertEquals(len(list(chunk_items([], len(members), 1))[0]), 0)
def test_raw_user_query(self):
invalid = ['', None, 12345, "select 'pretendemail'"]
for i in invalid:
items, count = execute_raw_user_query(i)
self.assertEquals(count, 0)
self.assertEquals(len(items), 0)
valid_sql = 'select * from %s' % user_table()
items, count = execute_raw_user_query(valid_sql)
self.assertEquals(count, 1)
self.assertEquals(len(items), 1)
class Segment(models.Model):
"""
A segment, as defined by a SQL query. Segments are designed to be stored in Redis and periodically refreshed.
"""
name = models.CharField(max_length=128)
slug = models.SlugField(max_length=256, null=True, blank=True, unique=True)
definition = models.TextField(
help_text="SQL query returning IDs of users in the segment.",
blank=True,
null=True)
priority = models.PositiveIntegerField(null=True, blank=True)
members_count = models.PositiveIntegerField(null=True,
blank=True,
default=0)
created_date = models.DateTimeField(auto_now_add=True)
updated_date = models.DateTimeField(null=True,
blank=True,
db_index=True,
auto_now=True)
recalculated_date = models.DateTimeField(null=True, blank=True)
helper = SegmentHelper()
############
# Public API
############
def has_member(self, user):
"""
Helper method. Return a bool indicating whether the user is a member of this segment.
"""
if not user.id:
return False
return self.helper.segment_has_member(self.id, user.id)
def add_member(self, user):
""" Helper method. Adds member to this segment. Returns a bool indicating the add status """
if not user.id:
return False
return self.helper.add_segment_membership(self.id, user.id)
@live_sql
def refresh(self):
members_count = self.helper.refresh_segment(self.id, self.definition)
Segment.objects.select_for_update().filter(id=self.id).update(
members_count=members_count, recalculated_date=timezone.now())
self.refresh_from_db()
def __len__(self):
"""Calling len() on a segment returns the number of members of that segment."""
return self.members_count
# A lot of code that interfaces with Django models expects model instances to be
# truthy, to distingiush them from `None`. Since we override `__len__`, `Segment`s
# with no members will be incorrectly treated as non-existent.
def __bool__(self):
return True
def __str__(self):
return self.name