Add context manager to FluentSender in fluent/sender.py

Kiseok Kim · Kiseok Kim · commit 3058514a5119 · 2017-03-14T04:40:37.000Z
This patch adds context manager with FluentSender and its test code.
diff --git a/fluent/sender.py b/fluent/sender.py
@@ -178,3 +178,9 @@ def _close(self):
         if self.socket:
             self.socket.close()
         self.socket = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, typ, value, traceback):
+        self.close()
diff --git a/tests/test_sender.py b/tests/test_sender.py
@@ -64,6 +64,18 @@ def test_simple(self):
         self.assertTrue(data[0][1])
         self.assertTrue(isinstance(data[0][1], int))
 
+    def test_decorator_simple(self):
+        with self._sender as sender:
+            sender.emit('foo', {'bar': 'baz'})
+        data = self.get_data()
+        eq = self.assertEqual
+        eq(1, len(data))
+        eq(3, len(data[0]))
+        eq('test.foo', data[0][0])
+        eq({'bar': 'baz'}, data[0][2])
+        self.assertTrue(data[0][1])
+        self.assertTrue(isinstance(data[0][1], int))
+
     def test_no_last_error_on_successful_emit(self):
         sender = self._sender
         sender.emit('foo', {'bar': 'baz'})
