Cleanup all test connections

Author: fayazkhan

sqlalchemy_mptt/tests/__init__.py

@@ -31,26 +31,33 @@
 
 """
 # standard library
-import os
+import contextlib
 import json
+import os
 import sys
+import typing
 import unittest
 
 # SQLAlchemy
 import sqlalchemy as sa
-from sqlalchemy import event, create_engine
+from sqlalchemy import create_engine, event
 from sqlalchemy.orm import sessionmaker
 
-# third-party
 from sqlalchemy_mptt import mptt_sessionmaker
+from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
 
-# local
-from .cases.get_tree import Tree
-from .cases.get_node import GetNodes
 from .cases.edit_node import Changes
+from .cases.get_node import GetNodes
+from .cases.get_tree import Tree
+from .cases.initialize import Initialize
 from .cases.integrity import DataIntegrity
 from .cases.move_node import MoveAfter, MoveBefore, MoveInside
-from .cases.initialize import Initialize
+
+if typing.TYPE_CHECKING:
+    BaseType = unittest.TestCase
+else:
+    BaseType = object
+DeclarativeBase = compat_layer.declarative_base()
 
 
 def failures_expected_on(*, sqlalchemy_versions=[], python_versions=[]):
@@ -73,6 +80,24 @@ def decorator(test_method):
     return decorator
 
 
+class DatabaseSetupMixin(BaseType):
+    base: DeclarativeBase  # type: ignore
+
+    def setUp(self):
+        with contextlib.suppress(AttributeError):
+            super().setUp()
+        self.engine: sa.engine.Engine = create_engine("sqlite:///:memory:")
+        Session = mptt_sessionmaker(sessionmaker(bind=self.engine))
+        self.session = Session()
+        self.base.metadata.create_all(self.engine)
+
+    def tearDown(self):
+        with contextlib.suppress(AttributeError):
+            super().tearDown()
+        self.session.close()
+        self.engine.dispose()
+
+
 class Fixtures(object):
     def __init__(self, session):
         self.session = session
@@ -97,6 +122,7 @@ class TreeTestingMixin(
     MoveInside,
     Tree,
     GetNodes,
+    DatabaseSetupMixin
 ):
     base = None
     model = None
@@ -116,10 +142,7 @@ def stop_query_counter(self):
         )
 
     def setUp(self):
-        self.engine = create_engine("sqlite:///:memory:")
-        Session = mptt_sessionmaker(sessionmaker(bind=self.engine))
-        self.session = Session()
-        self.base.metadata.create_all(self.engine)
+        super().setUp()
         self.fixture = Fixtures(self.session)
         self.fixture.add(
             self.model, os.path.join("fixtures", getattr(self, "fixtures", "tree.json"))
@@ -134,9 +157,6 @@ def setUp(self):
             self.model.tree_id,
         )
 
-    def tearDown(self):
-        self.base.metadata.drop_all(self.engine)
-
     def test_session_expire_for_move_after_to_new_tree(self):
         """
         https://github.com/uralbash/sqlalchemy_mptt/issues/33

sqlalchemy_mptt/tests/test_events.py

@@ -12,16 +12,12 @@
 
 import unittest
 
-from sqlalchemy import Column, Boolean, Integer, create_engine
+from sqlalchemy import Boolean, Column, Integer
 from sqlalchemy.event import contains
-from sqlalchemy.orm import sessionmaker
-
-from sqlalchemy_mptt import mptt_sessionmaker
 
 from sqlalchemy_mptt.mixins import BaseNestedSets
 from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
-from sqlalchemy_mptt.tests import TreeTestingMixin
-
+from sqlalchemy_mptt.tests import DatabaseSetupMixin, TreeTestingMixin
 
 Base = compat_layer.declarative_base()
 
@@ -156,40 +152,36 @@ def test_remove(self):
         tree_manager.register_events()
 
 
-class Tree0Id(unittest.TestCase):
+class Tree0Id(DatabaseSetupMixin, unittest.TestCase):
     """Test case where node id is provided and starts with 0
 
     See comments in https://github.com/uralbash/sqlalchemy_mptt/issues/57
     """
-    def test(self):
-        engine = create_engine('sqlite:///:memory:')
-        Session = mptt_sessionmaker(sessionmaker(bind=engine))
-        session = Session()
-        Base.metadata.create_all(engine)
 
+    base = Base
+
+    def test(self):
         root = Tree(id=0)
         child = Tree(id=1, parent_id=0)
 
-        session.add(root)
-        session.add(child)
-        session.commit()
+        self.session.add(root)
+        self.session.add(child)
+        self.session.commit()
 
         self.assertEqual(root.tree_id, 1)
         self.assertEqual(child.tree_id, 1)
 
 
-class InitialInsert(unittest.TestCase):
+class InitialInsert(DatabaseSetupMixin, unittest.TestCase):
     """Test case for initial insertion of node as specified in
     docs/initialize.rst
     """
+
+    base = Base
+
     def test_documented_initial_insert(self):
         from sqlalchemy_mptt import tree_manager
 
-        engine = create_engine('sqlite:///:memory:')
-        Session = mptt_sessionmaker(sessionmaker(bind=engine))
-        session = Session()
-        Base.metadata.create_all(engine)
-
         tree_manager.register_events(remove=True)  # Disable MPTT events
 
         _tree_id = 1
@@ -202,11 +194,11 @@ def test_documented_initial_insert(self):
                 right=0,
                 tree_id=_tree_id
             )
-            session.add(item)
-        session.commit()
+            self.session.add(item)
+        self.session.commit()
 
         tree_manager.register_events()  # enabled MPTT events back
         Tree.rebuild_tree(
-            session,
+            self.session,
             _tree_id
         )  # rebuild lft, rgt value automatically

sqlalchemy_mptt/tests/test_inheritance.py

@@ -1,12 +1,11 @@
 import unittest
 
 import sqlalchemy as sa
-from sqlalchemy.orm import sessionmaker
 
 from sqlalchemy_mptt.mixins import BaseNestedSets
 from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
-from sqlalchemy_mptt.tests import TreeTestingMixin, failures_expected_on
-
+from sqlalchemy_mptt.tests import (DatabaseSetupMixin, TreeTestingMixin,
+                                   failures_expected_on)
 
 Base = compat_layer.declarative_base()
 
@@ -45,16 +44,9 @@ class SpecializedTree(GenericTree):
     __table_args__ = tuple()
 
 
-class TestTree(unittest.TestCase):
-
-    def setUp(self):
-        self.engine = sa.create_engine('sqlite:///:memory:')
-        Session = sessionmaker(bind=self.engine)
-        self.session = Session()
-        Base.metadata.create_all(self.engine)
+class TestTree(DatabaseSetupMixin, unittest.TestCase):
 
-    def tearDown(self):
-        Base.metadata.drop_all(self.engine)
+    base = Base
 
     def test_create_generic(self):
         self.session.add(GenericTree(ppk=1))

sqlalchemy_mptt/tests/test_stateful.py

@@ -5,14 +5,16 @@
 #
 # Distributed under terms of the MIT license.
 """Test cases written using Hypothesis stateful testing framework."""
-from hypothesis import HealthCheck, settings, strategies as st
-from hypothesis.stateful import Bundle, RuleBasedStateMachine, consumes, invariant, rule
-from sqlalchemy import Column, Integer, Boolean, create_engine
-from sqlalchemy.orm import joinedload, sessionmaker
-
-from sqlalchemy_mptt import BaseNestedSets, mptt_sessionmaker
+from hypothesis import HealthCheck, settings
+from hypothesis import strategies as st
+from hypothesis.stateful import (Bundle, RuleBasedStateMachine, consumes,
+                                 invariant, rule)
+from sqlalchemy import Boolean, Column, Integer
+from sqlalchemy.orm import joinedload
+
+from sqlalchemy_mptt import BaseNestedSets
 from sqlalchemy_mptt.sqlalchemy_compat import compat_layer
-
+from sqlalchemy_mptt.tests import DatabaseSetupMixin
 
 Base = compat_layer.declarative_base()
 
@@ -27,15 +29,18 @@ def __repr__(self):
         return "<Node (%s)>" % self.id
 
 
-class TreeStateMachine(RuleBasedStateMachine):
+class TreeStateMachine(DatabaseSetupMixin, RuleBasedStateMachine):
     """A state machine with various possible actions and transitions for the Tree model."""
 
+    base = Base
+
     def __init__(self):
         super().__init__()
-        self.engine = create_engine("sqlite:///:memory:")
-        Session = mptt_sessionmaker(sessionmaker(bind=self.engine))
-        self.session = Session()
-        Base.metadata.create_all(self.engine)
+        self.setUp()
+
+    def teardown(self):
+        super().teardown()
+        self.tearDown()
 
     node = Bundle('node')