use contextmanager for layer and connection

Author: chaudum

pytest_crate/plugin.py

@@ -27,7 +27,14 @@ def __init__(self, name: str, version: str) -> None:
     def __repr__(self) -> str:
         return self.name
 
-    def start(self) -> None:
+    def __enter__(self):
+        self._start()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self._stop()
+
+    def _start(self) -> None:
         print(f"Starting CrateDB {self} ...")
         self.tmp = tempfile.mkdtemp()
         settings = {
@@ -41,7 +48,7 @@ def start(self) -> None:
         self.node.start()
         print(f"CrateDB {self} started")
 
-    def stop(self) -> None:
+    def _stop(self) -> None:
         print(f"Stopping CrateDB {self} ...")
         self.node.stop()
         shutil.rmtree(self.tmp, ignore_errors=True)
@@ -83,10 +90,8 @@ def crate_version(self, pytestconfig) -> Generator[str, None, None]:
     @pytest.fixture(scope="session")
     def crate_layer(self) -> CrateLayerFactoryGenerator:
         def layer_factory(name: str, version: str):
-            layer = CrateLayer(name, version)
-            layer.start()
-            yield layer
-            layer.stop()
+            with CrateLayer(name, version) as layer:
+                yield layer
         yield layer_factory
 
     @pytest.fixture(scope="session")
@@ -97,9 +102,8 @@ def crate(self, crate_layer, crate_version) -> CrateLayerGenerator:
 
     @pytest.fixture
     def crate_cursor(self, crate) -> Generator[Cursor, None, None]:
-        connection = connect(crate.dsn())
-        yield connection.cursor()
-        connection.close()
+        with connect(crate.dsn()) as connection:
+            yield connection.cursor()
 
     @pytest.fixture
     def crate_execute(self, crate_cursor) -> Generator[Callable, None, None]: