From 1155011b48a4cc99346615bdae0b484fd0ae3403 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= <kristof.toth@avatao.com>
Date: Sun, 23 Jun 2019 18:40:06 +0200
Subject: [PATCH] Make PipeIOServer a contextmanager

---
 pipe_io_server/pipe_io_server.py |  7 +++++++
 test_echo_server.py              | 24 ++++++------------------
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/pipe_io_server/pipe_io_server.py b/pipe_io_server/pipe_io_server.py
index 92eef98..02f12f6 100644
--- a/pipe_io_server/pipe_io_server.py
+++ b/pipe_io_server/pipe_io_server.py
@@ -77,3 +77,10 @@ class PipeIOServer(Thread):
 
     def wait(self):
         self._stop_event.wait()
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, type_, value, tb):
+        self.stop()
diff --git a/test_echo_server.py b/test_echo_server.py
index 56972db..7902602 100644
--- a/test_echo_server.py
+++ b/test_echo_server.py
@@ -5,7 +5,6 @@ from stat import S_ISFIFO
 from secrets import token_urlsafe
 from random import randint, getrandbits, uniform
 from threading import Thread
-from contextlib import contextmanager
 from json import dumps, loads
 
 import pytest
@@ -16,17 +15,14 @@ from pipe_io_server import PipeIOServer
 
 @pytest.fixture
 def io_pipes():
-    with pipe_io_server() as pipe_io:
+    with EchoPipeIOServer(*get_test_init_params()) as pipe_io:
         with IOPipes(pipe_io.in_pipe, pipe_io.out_pipe) as io_pipes:
             yield io_pipes
 
 
-@contextmanager
-def pipe_io_server(pipe_io_server_type=EchoPipeIOServer):
-    pipe_io_server = build_pipe_io_server(pipe_io_server_type)
-    pipe_io_server.start()
-    yield pipe_io_server
-    pipe_io_server.stop()
+def get_test_init_params():
+    here = dirname(realpath(__file__))
+    return join(here, 'in_pipe_tests'), join(here, 'out_pipe_tests')
 
 
 def raise_if_thread_blocks(thread_target_function, unblock_function):
@@ -38,14 +34,6 @@ def raise_if_thread_blocks(thread_target_function, unblock_function):
         raise RuntimeError('PipeIOServer failed to shut down!')
 
 
-def build_pipe_io_server(pipe_io_server_type=EchoPipeIOServer):
-    here = dirname(realpath(__file__))
-    return pipe_io_server_type(
-        join(here, 'in_pipe_tests'),
-        join(here, 'out_pipe_tests')
-    )
-
-
 class IOPipes:
     def __init__(self, in_pipe_path, out_pipe_path):
         self.in_pipe_path = in_pipe_path
@@ -79,7 +67,7 @@ def test_run_creates_pipes(io_pipes):
 
 
 def test_stop():
-    pipe_io = build_pipe_io_server()
+    pipe_io = EchoPipeIOServer(*get_test_init_params())
     def open_close_in_pipe():
         pipe_io.stop()
     raise_if_thread_blocks(pipe_io.run, open_close_in_pipe)
@@ -160,7 +148,7 @@ def test_json_io(io_pipes):
 
 
 def test_assign_message_handler():
-    pipe_io = build_pipe_io_server(PipeIOServer)
+    pipe_io = PipeIOServer(*get_test_init_params())
     pipe_io.handle_message = lambda msg: pipe_io.send_message(msg * 2)
     pipe_io.start()
     with IOPipes(pipe_io.in_pipe, pipe_io.out_pipe) as io_pipes: