Fix tests for per message sample_rate

2 files changed, 19 insertions, 26 deletions

diff --git a/tests/base/tests.py b/tests/base/tests.py
index 7c0a760..5644db0 100644
--- a/tests/base/tests.py
+++ b/tests/base/tests.py
@@ -557,36 +557,12 @@ class ClientTest(TestCase):
         self.assertEquals(len(self.client.events), 0)
 
     def test_sample_rate_per_message(self):
-        self.client.extra = {
-            'foo': 'bar',
-        }
         self.client.sample_rate = 1
-        self.client.captureMessage(message='test', extra={'sample_rate': 0.0})
+        self.client.captureMessage(message='test', sample_rate=0.0)
         self.assertEquals(len(self.client.events), 0)
 
         self.client.sample_rate = 0
-        self.client.captureMessage(message='test', extra={'sample_rate': 1.0})
-        self.assertEquals(len(self.client.events), 1)
-        event = self.client.events.pop(0)
-        if not PY2:
-            expected = {'sample_rate': 1.0, 'foo': "'bar'"}
-        else:
-            expected = {'sample_rate': 1.0, 'foo': "u'bar'"}
-        self.assertEquals(event['extra'], expected)
-
-    def test_sample_rate_per_message_is_resilient_to_bad_values(self):
-        self.client.sample_rate = 0
-
-        # sample_rate is not a number
-        self.client.captureMessage(message='test', extra={'sample_rate': 'foo'})
-        self.assertEquals(len(self.client.events), 0)
-
-        # sample_rate is not present
-        self.client.captureMessage(message='test', extra={'foo': '1.0'})
-        self.assertEquals(len(self.client.events), 0)
-
-        # sample_rate can be cast into a float
-        self.client.captureMessage(message='test', extra={'sample_rate': '1.0'})
+        self.client.captureMessage(message='test', sample_rate=1.0)
         self.assertEquals(len(self.client.events), 1)
 
     def test_transport_registration(self):
diff --git a/tests/handlers/logging/tests.py b/tests/handlers/logging/tests.py
index 163b8ab..8218b05 100644
--- a/tests/handlers/logging/tests.py
+++ b/tests/handlers/logging/tests.py
@@ -262,6 +262,23 @@ class LoggingIntegrationTest(TestCase):
         event = self.client.events.pop(0)
         assert event['server_name'] == 'foo'
 
+    def test_sample_rate(self):
+        record = self.make_record('Message', extra={'sample_rate': 0.0})
+        self.handler.emit(record)
+
+        self.assertEqual(len(self.client.events), 0)
+
+        record = self.make_record('Message', extra={'sample_rate': 1.0})
+        self.handler.emit(record)
+
+        self.assertEqual(len(self.client.events), 1)
+
+    def test_sample_rate_bad_values(self):
+        record = self.make_record('Message', extra={'sample_rate': 'foo'})
+        self.handler.emit(record)
+
+        self.assertEqual(len(self.client.events), 1)
+
 
 class LoggingHandlerTest(TestCase):
     def test_client_arg(self):