from contextlib import contextmanager

from django.conf import settings
from django.contrib.auth.models import User
from django.db import connection
from django.test import TestCase

from brabeion import badges
from brabeion.base import Badge, BadgeAwarded
from brabeion.tests.models import PlayerStat


class PointsBadge(Badge):
    slug = "points"
    levels = [
        "Bronze",
        "Silver",
        "Gold",
    ]
    events = [
        "points_awarded",
    ]
    multiple = False
    
    def award(self, **state):
        user = state["user"]
        points = user.stats.points
        if points > 10000:
            return BadgeAwarded(3)
        elif points > 7500:
            return BadgeAwarded(2)
        elif points > 5000:
            return BadgeAwarded(1)


badges.register(PointsBadge)


class BaseTestCase(TestCase):
    @contextmanager
    def assert_num_queries(self, n):
        current_debug = settings.DEBUG
        settings.DEBUG = True
        current = len(connection.queries)
        yield
        self.assertEqual(current+n, len(connection.queries), connection.queries[current:])
        settings.DEBUG = current_debug

class BadgesTests(BaseTestCase):
    def test_award(self):
        u = User.objects.create_user("Lars Bak", "lars@hotspot.com", "x864lyfe")
        PlayerStat.objects.create(user=u)
        badges.possibly_award_badge("points_awarded", user=u)
        self.assertEqual(u.badges_earned.count(), 0)

        u.stats.points += 5001
        u.stats.save()
        badges.possibly_award_badge("points_awarded", user=u)
        self.assertEqual(u.badges_earned.count(), 1)
        self.assertEqual(u.badges_earned.all()[0].badge.name, "Bronze")

        badges.possibly_award_badge("points_awarded", user=u)
        self.assertEqual(u.badges_earned.count(), 1)

        u.stats.points += 2500
        badges.possibly_award_badge("points_awarded", user=u)
        self.assertEqual(u.badges_earned.count(), 2)
    
    def test_lazy_user(self):
        u = User.objects.create_user("Lars Bak", "lars@hotspot.com", "x864lyfe")
        PlayerStat.objects.create(user=u, points=5001)
        badges.possibly_award_badge("points_awarded", user=u)
        self.assertEqual(u.badges_earned.count(), 1)
        
        with self.assert_num_queries(1):
            u.badges_earned.get().badge

@@ -1,43 +1,4 @@
-from contextlib import contextmanager%0A%0A
 from
@@ -799,28 +799,8 @@
 e):%0A
-    @contextmanager%0A
     
@@ -829,16 +829,22 @@
 (self, n
+, func
 ):%0A     
@@ -961,13 +961,14 @@
     
-yield
+func()
 %0A   
@@ -2204,13 +2204,8 @@
     
-with 
 self
@@ -2229,22 +2229,17 @@
 es(1
-):%0A           
+, lambda:
  u.b
@@ -2262,9 +2262,10 @@
 ().badge
+)
 %0A

Work with python2.4, losing the context manager hotness is sad.