From 16bf4e945bb32d121940ebe7ddc19b96f9e823ab Mon Sep 17 00:00:00 2001
From: Chenxiong Qi <cqi@redhat.com>
Date: Sun, 30 Jun 2019 11:05:57 +0800
Subject: [PATCH] Reuse ModuleBuild.get_by_id

Signed-off-by: Chenxiong Qi <cqi@redhat.com>
---
 tests/__init__.py                        |  4 +-
 tests/test_build/test_build.py           |  2 +-
 tests/test_builder/test_koji.py          | 18 ++++---
 tests/test_models/test_models.py         | 17 +++---
 tests/test_resolver/test_db.py           | 39 +++++++-------
 tests/test_scheduler/test_module_init.py | 25 ++++-----
 tests/test_scheduler/test_poller.py      | 69 +++++++++++++-----------
 tests/test_scheduler/test_repo_done.py   | 10 ++--
 tests/test_utils/test_utils.py           | 63 +++++++++++-----------
 tests/test_views/test_views.py           |  8 +--
 10 files changed, 132 insertions(+), 123 deletions(-)

diff --git a/tests/__init__.py b/tests/__init__.py
index caa2ade3..29e0bfea 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -368,7 +368,7 @@ def scheduler_init_data(db_session, tangerine_state=None, scratch=False):
     db_session.add(module_build)
     db_session.commit()
 
-    platform_br = db_session.query(module_build_service.models.ModuleBuild).get(1)
+    platform_br = module_build_service.models.ModuleBuild.get_by_id(db_session, 1)
     module_build.buildrequires.append(platform_br)
 
     arch = db_session.query(module_build_service.models.ModuleArch).get(1)
@@ -477,7 +477,7 @@ def reuse_component_init_data():
     db.session.commit()
     db.session.refresh(build_one)
 
-    platform_br = module_build_service.models.ModuleBuild.query.get(1)
+    platform_br = module_build_service.models.ModuleBuild.get_by_id(db.session, 1)
     build_one.buildrequires.append(platform_br)
 
     arch = module_build_service.models.ModuleArch.query.get(1)
diff --git a/tests/test_build/test_build.py b/tests/test_build/test_build.py
index 0fcb028e..925db2c3 100644
--- a/tests/test_build/test_build.py
+++ b/tests/test_build/test_build.py
@@ -1650,7 +1650,7 @@ class TestBuild(BaseTestBuild):
 
         # Simulate a random repo regen message that MBS didn't expect
         cleanup_moksha()
-        module = db_session.query(models.ModuleBuild).get(module_build_id)
+        module = models.ModuleBuild.get_by_id(db_session, module_build_id)
         msgs = [
             module_build_service.messaging.KojiRepoChange(
                 msg_id="a faked internal message", repo_tag=module.koji_tag + "-build"
diff --git a/tests/test_builder/test_koji.py b/tests/test_builder/test_koji.py
index d21a9c07..002f303d 100644
--- a/tests/test_builder/test_koji.py
+++ b/tests/test_builder/test_koji.py
@@ -136,7 +136,7 @@ class TestKojiBuilder:
             "/module-base-runtime-0.25-9/latest/x86_64"
         )
 
-    def test_recover_orphaned_artifact_when_tagged(self):
+    def test_recover_orphaned_artifact_when_tagged(self, db_session):
         """ Test recover_orphaned_artifact when the artifact is found and tagged in both tags
         """
         builder = FakeKojiModuleBuilder(
@@ -154,7 +154,7 @@ class TestKojiBuilder:
         build_tagged = [{"nvr": "foo-1.0-1.module+e0095747", "task_id": 12345, "build_id": 91}]
         dest_tagged = [{"nvr": "foo-1.0-1.module+e0095747", "task_id": 12345, "build_id": 91}]
         builder.koji_session.listTagged.side_effect = [build_tagged, dest_tagged]
-        module_build = module_build_service.models.ModuleBuild.query.get(4)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 4)
         component_build = module_build.component_builds[0]
         component_build.task_id = None
         component_build.state = None
@@ -181,7 +181,7 @@ class TestKojiBuilder:
         assert component_build.state_reason == "Found existing build"
         assert builder.koji_session.tagBuild.call_count == 0
 
-    def test_recover_orphaned_artifact_when_untagged(self):
+    def test_recover_orphaned_artifact_when_untagged(self, db_session):
         """ Tests recover_orphaned_artifact when the build is found but untagged
         """
         builder = FakeKojiModuleBuilder(
@@ -203,7 +203,7 @@ class TestKojiBuilder:
         builder.koji_session.untaggedBuilds.return_value = untagged
         build_info = {"nvr": "foo-1.0-1.{0}".format(dist_tag), "task_id": 12345, "build_id": 91}
         builder.koji_session.getBuild.return_value = build_info
-        module_build = module_build_service.models.ModuleBuild.query.get(4)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 4)
         component_build = module_build.component_builds[0]
         component_build.task_id = None
         component_build.nvr = None
@@ -224,7 +224,7 @@ class TestKojiBuilder:
         assert component_build.state_reason == "Found existing build"
         builder.koji_session.tagBuild.assert_called_once_with(2, "foo-1.0-1.{0}".format(dist_tag))
 
-    def test_recover_orphaned_artifact_when_nothing_exists(self):
+    def test_recover_orphaned_artifact_when_nothing_exists(self, db_session):
         """ Test recover_orphaned_artifact when the build is not found
         """
         builder = FakeKojiModuleBuilder(
@@ -243,7 +243,7 @@ class TestKojiBuilder:
         builder.koji_session.listTagged.return_value = tagged
         untagged = [{"nvr": "foo-1.0-1.nope", "release": "nope"}]
         builder.koji_session.untaggedBuilds.return_value = untagged
-        module_build = module_build_service.models.ModuleBuild.query.get(4)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 4)
         component_build = module_build.component_builds[0]
         component_build.task_id = None
         component_build.nvr = None
@@ -729,7 +729,9 @@ class TestKojiBuilder:
         ),
     )
     @patch("module_build_service.builder.KojiModuleBuilder.KojiClientSession")
-    def test_get_filtered_rpms_on_self_dep(self, ClientSession, br_filtered_rpms, expected):
+    def test_get_filtered_rpms_on_self_dep(
+        self, ClientSession, br_filtered_rpms, expected, db_session
+    ):
         session = ClientSession.return_value
         session.listTaggedRPMS.return_value = (
             [
@@ -774,7 +776,7 @@ class TestKojiBuilder:
             ],
         )
         reuse_component_init_data()
-        current_module = module_build_service.models.ModuleBuild.query.get(3)
+        current_module = module_build_service.models.ModuleBuild.get_by_id(db_session, 3)
         rv = KojiModuleBuilder._get_filtered_rpms_on_self_dep(current_module, br_filtered_rpms)
         assert set(rv) == set(expected)
         session.assert_not_called()
diff --git a/tests/test_models/test_models.py b/tests/test_models/test_models.py
index 0ab3d54a..c3886c06 100644
--- a/tests/test_models/test_models.py
+++ b/tests/test_models/test_models.py
@@ -77,20 +77,19 @@ class TestModels:
         assert build.runtime_context == "bbc84c7b817ab3dd54916c0bcd6c6bdf512f7f9c"
         assert build.context == "3ee22b28"
 
-    def test_siblings_property(self):
+    def test_siblings_property(self, db_session):
         """ Tests that the siblings property returns the ID of all modules with
         the same name:stream:version
         """
         clean_database()
         mmd = load_mmd(read_staged_data("formatted_testmodule"))
-        with make_session(conf) as session:
-            for i in range(3):
-                build = module_build_from_modulemd(mmd_to_str(mmd))
-                build.build_context = "f6e2aeec7576196241b9afa0b6b22acf2b6873d" + str(i)
-                build.runtime_context = "bbc84c7b817ab3dd54916c0bcd6c6bdf512f7f9c" + str(i)
-                session.add(build)
-        session.commit()
-        build_one = ModuleBuild.query.get(2)
+        for i in range(3):
+            build = module_build_from_modulemd(mmd_to_str(mmd))
+            build.build_context = "f6e2aeec7576196241b9afa0b6b22acf2b6873d" + str(i)
+            build.runtime_context = "bbc84c7b817ab3dd54916c0bcd6c6bdf512f7f9c" + str(i)
+            db_session.add(build)
+        db_session.commit()
+        build_one = ModuleBuild.get_by_id(db_session, 2)
         assert build_one.siblings == [3, 4]
 
     @pytest.mark.parametrize(
diff --git a/tests/test_resolver/test_db.py b/tests/test_resolver/test_db.py
index 6ad0bbea..5180b4ea 100644
--- a/tests/test_resolver/test_db.py
+++ b/tests/test_resolver/test_db.py
@@ -90,15 +90,15 @@ class TestDBModule:
             assert nsvcs == set(["platform:f29.1.0:3:00000000"])
 
     @pytest.mark.parametrize("empty_buildrequires", [False, True])
-    def test_get_module_build_dependencies(self, empty_buildrequires):
+    def test_get_module_build_dependencies(self, empty_buildrequires, db_session):
         """
         Tests that the buildrequires of testmodule are returned
         """
         expected = set(["module-f28-build"])
-        module = models.ModuleBuild.query.get(2)
+        module = models.ModuleBuild.get_by_id(db_session, 2)
         if empty_buildrequires:
             expected = set()
-            module = models.ModuleBuild.query.get(2)
+            module = models.ModuleBuild.get_by_id(db_session, 2)
             mmd = module.mmd()
             # Wipe out the dependencies
             mmd.clear_dependencies()
@@ -106,19 +106,18 @@ class TestDBModule:
             xmd["mbs"]["buildrequires"] = {}
             mmd.set_xmd(xmd)
             module.modulemd = mmd_to_str(mmd)
-            db.session.add(module)
-            db.session.commit()
+            db_session.commit()
         resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="db")
         result = resolver.get_module_build_dependencies(
             "testmodule", "master", "20170109091357", "78e4a6fd").keys()
         assert set(result) == expected
 
-    def test_get_module_build_dependencies_recursive(self):
+    def test_get_module_build_dependencies_recursive(self, db_session):
         """
         Tests that the buildrequires are returned when it is two layers deep
         """
         # Add testmodule2 that requires testmodule
-        module = models.ModuleBuild.query.get(3)
+        module = models.ModuleBuild.get_by_id(db_session, 3)
         mmd = module.mmd()
         # Rename the module
         mmd = mmd.copy("testmodule2")
@@ -139,8 +138,7 @@ class TestDBModule:
         module.version = str(mmd.get_version())
         module.koji_tag = "module-ae2adf69caf0e1b6"
 
-        db.session.add(module)
-        db.session.commit()
+        db_session.commit()
 
         resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="db")
         result = resolver.get_module_build_dependencies(
@@ -171,8 +169,8 @@ class TestDBModule:
             expected = [os.path.join(local_path, "module-parent-master-20170816080815/results")]
             assert set(result) == set(expected)
 
-    def test_resolve_requires(self):
-        build = models.ModuleBuild.query.get(2)
+    def test_resolve_requires(self, db_session):
+        build = models.ModuleBuild.get_by_id(db_session, 2)
         resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="db")
         result = resolver.resolve_requires(
             [":".join([build.name, build.stream, build.version, build.context])]
@@ -188,11 +186,11 @@ class TestDBModule:
             }
         }
 
-    def test_resolve_profiles(self):
+    def test_resolve_profiles(self, db_session):
         """
         Tests that the profiles get resolved recursively
         """
-        mmd = models.ModuleBuild.query.get(2).mmd()
+        mmd = models.ModuleBuild.get_by_id(db_session, 2).mmd()
         resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="db")
         result = resolver.resolve_profiles(mmd, ("buildroot", "srpm-buildroot"))
         expected = {
@@ -242,17 +240,16 @@ class TestDBModule:
         new_callable=PropertyMock,
         return_value=tests.staged_data_filename("local_builds")
     )
-    def test_resolve_profiles_local_module(self, local_builds, conf_system):
+    def test_resolve_profiles_local_module(self, local_builds, conf_system, db_session):
         """
         Test that profiles get resolved recursively on local builds
         """
-        with app.app_context():
-            utils.load_local_builds(["platform"])
-            mmd = models.ModuleBuild.query.get(2).mmd()
-            resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="mbs")
-            result = resolver.resolve_profiles(mmd, ("buildroot", "srpm-buildroot"))
-            expected = {"buildroot": set(["foo"]), "srpm-buildroot": set(["bar"])}
-            assert result == expected
+        utils.load_local_builds(["platform"])
+        mmd = models.ModuleBuild.get_by_id(db_session, 2).mmd()
+        resolver = mbs_resolver.GenericResolver.create(tests.conf, backend="mbs")
+        result = resolver.resolve_profiles(mmd, ("buildroot", "srpm-buildroot"))
+        expected = {"buildroot": set(["foo"]), "srpm-buildroot": set(["bar"])}
+        assert result == expected
 
     def test_get_latest_with_virtual_stream(self):
         tests.init_data(1, multiple_stream_versions=True)
diff --git a/tests/test_scheduler/test_module_init.py b/tests/test_scheduler/test_module_init.py
index 8f0f2d80..affa13b3 100644
--- a/tests/test_scheduler/test_module_init.py
+++ b/tests/test_scheduler/test_module_init.py
@@ -27,7 +27,7 @@ from tests import conf, clean_database, read_staged_data
 from tests.test_views.test_views import FakeSCM
 import module_build_service.messaging
 import module_build_service.scheduler.handlers.modules
-from module_build_service import build_logs, db
+from module_build_service import build_logs
 from module_build_service.models import make_session, ModuleBuild, ComponentBuild
 from module_build_service.utils.general import mmd_to_str, load_mmd
 
@@ -59,7 +59,7 @@ class TestModuleInit:
     @patch("module_build_service.scm.SCM")
     @patch("module_build_service.scheduler.handlers.modules.handle_stream_collision_modules")
     @patch("module_build_service.utils.submit.get_build_arches", return_value=["x86_64"])
-    def test_init_basic(self, get_build_arches, rscm, mocked_scm, built_rpms):
+    def init_basic(self, db_session, get_build_arches, rscm, mocked_scm, built_rpms):
         FakeSCM(
             mocked_scm,
             "testmodule",
@@ -76,22 +76,23 @@ class TestModuleInit:
             "x-debuginfo-0:2.5.48-3.el8+1308+551bfa71",
         ]
 
-        platform_build = ModuleBuild.query.get(1)
+        platform_build = ModuleBuild.get_by_id(db_session, 1)
         mmd = platform_build.mmd()
         for rpm in mmd.get_rpm_filters():
             mmd.remove_rpm_filter(rpm)
         mmd.add_rpm_filter("foo")
         mmd.add_rpm_filter("bar")
+
         platform_build.modulemd = mmd_to_str(mmd)
-        db.session.commit()
+        db_session.commit()
 
         msg = module_build_service.messaging.MBSModule(
             msg_id=None, module_build_id=2, module_build_state="init"
         )
 
-        with make_session(conf) as session:
-            self.fn(config=conf, session=session, msg=msg)
-        build = ModuleBuild.query.filter_by(id=2).one()
+        self.fn(config=conf, session=db_session, msg=msg)
+
+        build = ModuleBuild.get_by_id(db_session, 2)
         # Make sure the module entered the wait state
         assert build.state == 1, build.state
         # Make sure format_mmd was run properly
@@ -102,15 +103,15 @@ class TestModuleInit:
         ]
         return build
 
-    def test_init_called_twice(self):
-        build = self.test_init_basic()
+    def test_init_called_twice(self, db_session):
+        build = self.init_basic(db_session)
         old_component_builds = len(build.component_builds)
         old_mmd = load_mmd(build.modulemd)
 
         build.state = 4
-        db.session.commit()
-        build = self.test_init_basic()
-        db.session.refresh(build)
+        db_session.commit()
+        build = self.init_basic(db_session)
+        db_session.refresh(build)
 
         assert build.state == 1
         assert old_component_builds == len(build.component_builds)
diff --git a/tests/test_scheduler/test_poller.py b/tests/test_scheduler/test_poller.py
index fae0f7b6..0951a045 100644
--- a/tests/test_scheduler/test_poller.py
+++ b/tests/test_scheduler/test_poller.py
@@ -442,28 +442,29 @@ class TestPoller:
         # Ensure we did *not* process any of the non-waiting builds.
         assert consumer.incoming.qsize() == 0
 
-    def test_cleanup_stale_failed_builds(self, create_builder, global_consumer, dbg):
+    def test_cleanup_stale_failed_builds(self, create_builder, global_consumer, dbg, db_session):
         """ Test that one of the two module builds gets to the garbage state when running
         cleanup_stale_failed_builds.
         """
         builder = mock.MagicMock()
         create_builder.return_value = builder
-        module_build_one = models.ModuleBuild.query.get(2)
-        module_build_two = models.ModuleBuild.query.get(3)
+
+        module_build_one = models.ModuleBuild.get_by_id(db_session, 2)
         module_build_one.state = models.BUILD_STATES["failed"]
         module_build_one.time_modified = datetime.utcnow() - timedelta(
             days=conf.cleanup_failed_builds_time + 1)
+
+        module_build_two = models.ModuleBuild.get_by_id(db_session, 3)
         module_build_two.time_modified = datetime.utcnow()
         module_build_two.state = models.BUILD_STATES["failed"]
-        failed_component = models.ComponentBuild.query.filter_by(
+
+        failed_component = db_session.query(models.ComponentBuild).filter_by(
             package="tangerine", module_id=3).one()
         failed_component.state = koji.BUILD_STATES["FAILED"]
         failed_component.tagged = False
         failed_component.tagged_in_final = False
-        db.session.add(failed_component)
-        db.session.add(module_build_one)
-        db.session.add(module_build_two)
-        db.session.commit()
+
+        db_session.commit()
 
         consumer = mock.MagicMock()
         consumer.incoming = queue.Queue()
@@ -473,8 +474,8 @@ class TestPoller:
 
         # Ensure the queue is empty before we start
         assert consumer.incoming.qsize() == 0
-        poller.cleanup_stale_failed_builds(conf, db.session)
-        db.session.refresh(module_build_two)
+        poller.cleanup_stale_failed_builds(conf, db_session)
+        db_session.refresh(module_build_two)
         # Make sure module_build_one was transitioned to garbage
         assert module_build_one.state == models.BUILD_STATES["garbage"]
         state_reason = (
@@ -496,25 +497,27 @@ class TestPoller:
             "module-build-macros-0.1-1.module+0+d027b723",
         ])
 
-    def test_cleanup_stale_failed_builds_no_components(self, create_builder, global_consumer, dbg):
+    def test_cleanup_stale_failed_builds_no_components(
+        self, create_builder, global_consumer, dbg, db_session
+    ):
         """ Test that a module build without any components built gets to the garbage state when
         running cleanup_stale_failed_builds.
         """
-        module_build_one = models.ModuleBuild.query.get(1)
-        module_build_two = models.ModuleBuild.query.get(2)
+        module_build_one = models.ModuleBuild.get_by_id(db_session, 1)
         module_build_one.state = models.BUILD_STATES["failed"]
         module_build_one.time_modified = datetime.utcnow()
+
+        module_build_two = models.ModuleBuild.get_by_id(db_session, 2)
         module_build_two.state = models.BUILD_STATES["failed"]
         module_build_two.time_modified = datetime.utcnow() - timedelta(
             days=conf.cleanup_failed_builds_time + 1)
         module_build_two.koji_tag = None
         module_build_two.cg_build_koji_tag = None
+
         for c in module_build_two.component_builds:
             c.state = None
-            db.session.add(c)
-        db.session.add(module_build_one)
-        db.session.add(module_build_two)
-        db.session.commit()
+
+        db_session.commit()
 
         consumer = mock.MagicMock()
         consumer.incoming = queue.Queue()
@@ -524,8 +527,8 @@ class TestPoller:
 
         # Ensure the queue is empty before we start
         assert consumer.incoming.qsize() == 0
-        poller.cleanup_stale_failed_builds(conf, db.session)
-        db.session.refresh(module_build_two)
+        poller.cleanup_stale_failed_builds(conf, db_session)
+        db_session.refresh(module_build_two)
         # Make sure module_build_two was transitioned to garbage
         assert module_build_two.state == models.BUILD_STATES["garbage"]
         state_reason = (
@@ -541,24 +544,26 @@ class TestPoller:
     @pytest.mark.parametrize(
         "test_state", [models.BUILD_STATES[state] for state in conf.cleanup_stuck_builds_states]
     )
-    def test_cancel_stuck_module_builds(self, create_builder, global_consumer, dbg, test_state):
+    def test_cancel_stuck_module_builds(
+        self, create_builder, global_consumer, dbg, test_state, db_session
+    ):
 
-        module_build1 = models.ModuleBuild.query.get(1)
+        module_build1 = models.ModuleBuild.get_by_id(db_session, 1)
         module_build1.state = test_state
         under_thresh = conf.cleanup_stuck_builds_time - 1
         module_build1.time_modified = datetime.utcnow() - timedelta(
             days=under_thresh, hours=23, minutes=59)
 
-        module_build2 = models.ModuleBuild.query.get(2)
+        module_build2 = models.ModuleBuild.get_by_id(db_session, 2)
         module_build2.state = test_state
         module_build2.time_modified = datetime.utcnow() - timedelta(
             days=conf.cleanup_stuck_builds_time)
 
-        module_build2 = models.ModuleBuild.query.get(3)
+        module_build2 = models.ModuleBuild.get_by_id(db_session, 3)
         module_build2.state = test_state
         module_build2.time_modified = datetime.utcnow()
 
-        db.session.commit()
+        db_session.commit()
 
         consumer = mock.MagicMock()
         consumer.incoming = queue.Queue()
@@ -568,9 +573,9 @@ class TestPoller:
 
         assert consumer.incoming.qsize() == 0
 
-        poller.cancel_stuck_module_builds(conf, db.session)
+        poller.cancel_stuck_module_builds(conf, db_session)
 
-        module = models.ModuleBuild.query.filter_by(state=4).all()
+        module = db_session.query(models.ModuleBuild).filter_by(state=4).all()
         assert len(module) == 1
         assert module[0].id == 2
 
@@ -634,18 +639,20 @@ class TestPoller:
 
     @pytest.mark.parametrize("greenwave_result", [True, False])
     @patch("module_build_service.utils.greenwave.Greenwave.check_gating")
-    def test_poll_greenwave(self, mock_gw, create_builder, global_consumer, dbg, greenwave_result):
+    def test_poll_greenwave(
+        self, mock_gw, create_builder, global_consumer, dbg, greenwave_result, db_session
+    ):
 
-        module_build1 = models.ModuleBuild.query.get(1)
+        module_build1 = models.ModuleBuild.get_by_id(db_session, 1)
         module_build1.state = models.BUILD_STATES["ready"]
 
-        module_build2 = models.ModuleBuild.query.get(2)
+        module_build2 = models.ModuleBuild.get_by_id(db_session, 2)
         module_build2.state = models.BUILD_STATES["done"]
 
-        module_build2 = models.ModuleBuild.query.get(3)
+        module_build2 = models.ModuleBuild.get_by_id(db_session, 3)
         module_build2.state = models.BUILD_STATES["init"]
 
-        db.session.commit()
+        db_session.commit()
 
         consumer = mock.MagicMock()
         consumer.incoming = queue.Queue()
diff --git a/tests/test_scheduler/test_repo_done.py b/tests/test_scheduler/test_repo_done.py
index 5dda406a..1058bc80 100644
--- a/tests/test_scheduler/test_repo_done.py
+++ b/tests/test_scheduler/test_repo_done.py
@@ -122,15 +122,15 @@ class TestRepoDone:
 
         # Ensure the time_completed is None, so we can test it is set to
         # some date once the build is finalized.
-        module_build = module_build_service.models.ModuleBuild.query.get(2)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 2)
         module_build.time_completed = None
-        db.session.commit()
+        db_session.commit()
 
         def mocked_finalizer(succeeded=None):
             # Check that the time_completed is set in the time when
             # finalizer is called.
             assert succeeded is True
-            module_build = db_session.query(module_build_service.models.ModuleBuild).get(2)
+            module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 2)
             assert module_build.time_completed is not None
 
         finalizer.side_effect = mocked_finalizer
@@ -207,7 +207,7 @@ class TestRepoDone:
         mock_log_info.assert_called_with(
             "Ignoring repo regen, because not all components are tagged."
         )
-        module_build = module_build_service.models.ModuleBuild.query.get(2)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 2)
         # Make sure the module build didn't transition since all the components weren't tagged
         assert module_build.state == module_build_service.models.BUILD_STATES["build"]
 
@@ -242,6 +242,6 @@ class TestRepoDone:
         msg = module_build_service.messaging.KojiRepoChange(
             "some_msg_id", "module-testmodule-master-20170109091357-7c29193d-build")
         module_build_service.scheduler.handlers.repos.done(config=conf, session=db_session, msg=msg)
-        module_build = module_build_service.models.ModuleBuild.query.get(2)
+        module_build = module_build_service.models.ModuleBuild.get_by_id(db_session, 2)
 
         assert module_build.state == module_build_service.models.BUILD_STATES["failed"]
diff --git a/tests/test_utils/test_utils.py b/tests/test_utils/test_utils.py
index 82c20276..52274c8b 100644
--- a/tests/test_utils/test_utils.py
+++ b/tests/test_utils/test_utils.py
@@ -436,18 +436,20 @@ class TestUtils:
         else:
             module_build_service.utils.import_mmd(db.session, mmd)
 
-    def test_get_rpm_release_mse(self):
+    def test_get_rpm_release_mse(self, db_session):
         init_data(contexts=True)
-        build_one = models.ModuleBuild.query.get(2)
-        build_two = models.ModuleBuild.query.get(3)
+
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         release_one = module_build_service.utils.get_rpm_release(build_one)
-        release_two = module_build_service.utils.get_rpm_release(build_two)
         assert release_one == "module+2+b8645bbb"
+
+        build_two = models.ModuleBuild.get_by_id(db_session, 3)
+        release_two = module_build_service.utils.get_rpm_release(build_two)
         assert release_two == "module+2+17e35784"
 
     def test_get_rpm_release_platform_stream(self, db_session):
         scheduler_init_data(db_session, 1)
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         release = module_build_service.utils.get_rpm_release(build_one)
         assert release == "module+f28+2+814cfa39"
 
@@ -468,7 +470,7 @@ class TestUtils:
         db_session.add(platform)
         db_session.commit()
 
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         release = module_build_service.utils.get_rpm_release(build_one)
         assert release == "module+fedora28+2+814cfa39"
 
@@ -486,7 +488,7 @@ class TestUtils:
         metadata_mmd = load_mmd(read_staged_data("build_metadata_module"))
         module_build_service.utils.import_mmd(db_session, metadata_mmd)
 
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         mmd = build_one.mmd()
         deps = mmd.get_dependencies()[0]
         deps.add_buildtime_stream("build", "product1.2")
@@ -506,18 +508,20 @@ class TestUtils:
         release = module_build_service.utils.get_rpm_release(build_one)
         assert release == "module+product12+2+814cfa39"
 
-    def test_get_rpm_release_mse_scratch(self):
+    def test_get_rpm_release_mse_scratch(self, db_session):
         init_data(contexts=True, scratch=True)
-        build_one = models.ModuleBuild.query.get(2)
-        build_two = models.ModuleBuild.query.get(3)
+
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         release_one = module_build_service.utils.get_rpm_release(build_one)
-        release_two = module_build_service.utils.get_rpm_release(build_two)
         assert release_one == "scrmod+2+b8645bbb"
+
+        build_two = models.ModuleBuild.get_by_id(db_session, 3)
+        release_two = module_build_service.utils.get_rpm_release(build_two)
         assert release_two == "scrmod+2+17e35784"
 
     def test_get_rpm_release_platform_stream_scratch(self, db_session):
         scheduler_init_data(db_session, 1, scratch=True)
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         release = module_build_service.utils.get_rpm_release(build_one)
         assert release == "scrmod+f28+2+814cfa39"
 
@@ -525,7 +529,7 @@ class TestUtils:
     def test_record_module_build_arches(self, get_build_arches, db_session):
         get_build_arches.return_value = ["x86_64", "i686"]
         scheduler_init_data(db_session, 1)
-        build = db_session.query(models.ModuleBuild).get(2)
+        build = models.ModuleBuild.get_by_id(db_session, 2)
         build.arches = []
         module_build_service.utils.record_module_build_arches(build.mmd(), build, db_session)
 
@@ -592,7 +596,7 @@ class TestUtils:
         mmd_xmd = mmd.get_xmd()
         assert mmd_xmd == xmd
 
-    def test_get_reusable_component_shared_userspace_ordering(self):
+    def test_get_reusable_component_shared_userspace_ordering(self, db_session):
         """
         For modules with lot of components per batch, there is big chance that
         the database will return them in different order than what we have for
@@ -600,7 +604,7 @@ class TestUtils:
         reuse the components.
         """
         reuse_shared_userspace_init_data()
-        new_module = models.ModuleBuild.query.get(3)
+        new_module = models.ModuleBuild.get_by_id(db_session, 3)
         rv = module_build_service.utils.get_reusable_component(db.session, new_module, "llvm")
         assert rv.package == "llvm"
 
@@ -830,24 +834,23 @@ class TestUtils:
 
     @patch("module_build_service.scm.SCM")
     @patch("module_build_service.utils.submit.ThreadPool")
-    def test_format_mmd_update_time_modified(self, tp, mocked_scm):
-        with app.app_context():
-            init_data()
-            build = models.ModuleBuild.query.get(2)
+    def test_format_mmd_update_time_modified(self, tp, mocked_scm, db_session):
+        init_data()
+        build = models.ModuleBuild.get_by_id(db_session, 2)
 
-            async_result = mock.MagicMock()
-            async_result.ready.side_effect = [False, False, False, True]
-            tp.return_value.map_async.return_value = async_result
+        async_result = mock.MagicMock()
+        async_result.ready.side_effect = [False, False, False, True]
+        tp.return_value.map_async.return_value = async_result
 
-            test_datetime = datetime(2019, 2, 14, 11, 11, 45, 42968)
+        test_datetime = datetime(2019, 2, 14, 11, 11, 45, 42968)
 
-            mmd = load_mmd(read_staged_data("testmodule"))
+        mmd = load_mmd(read_staged_data("testmodule"))
 
-            with patch("module_build_service.utils.submit.datetime") as dt:
-                dt.utcnow.return_value = test_datetime
-                module_build_service.utils.format_mmd(mmd, None, build, db.session)
+        with patch("module_build_service.utils.submit.datetime") as dt:
+            dt.utcnow.return_value = test_datetime
+            module_build_service.utils.format_mmd(mmd, None, build, db_session)
 
-            assert build.time_modified == test_datetime
+        assert build.time_modified == test_datetime
 
     def test_generate_koji_tag_in_nsvc_format(self):
         name, stream, version, context = ("testmodule", "master", "20170816080815", "37c6c57")
@@ -893,13 +896,13 @@ class TestUtils:
 
     def test_get_prefixed_version_f28(self, db_session):
         scheduler_init_data(db_session, 1)
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         v = module_build_service.utils.submit.get_prefixed_version(build_one.mmd())
         assert v == 2820180205135154
 
     def test_get_prefixed_version_fl701(self, db_session):
         scheduler_init_data(db_session, 1)
-        build_one = db_session.query(models.ModuleBuild).get(2)
+        build_one = models.ModuleBuild.get_by_id(db_session, 2)
         mmd = build_one.mmd()
         xmd = mmd.get_xmd()
         xmd["mbs"]["buildrequires"]["platform"]["stream"] = "fl7.0.1-beta"
diff --git a/tests/test_views/test_views.py b/tests/test_views/test_views.py
index cffe8fe5..41836a24 100644
--- a/tests/test_views/test_views.py
+++ b/tests/test_views/test_views.py
@@ -758,7 +758,7 @@ class TestViews:
 
     def test_query_builds_order_by_multiple(self):
         init_data(data_size=1, multiple_stream_versions=True)
-        platform_f28 = db.session.query(module_build_service.models.ModuleBuild).get(1)
+        platform_f28 = module_build_service.models.ModuleBuild.get_by_id(db.session, 1)
         platform_f28.version = "150"
         db.session.add(platform_f28)
         db.session.commit()
@@ -945,7 +945,7 @@ class TestViews:
         module_build_service.utils.load_mmd(data["modulemd"])
 
         # Make sure the buildrequires entry was created
-        module = ModuleBuild.query.get(8)
+        module = ModuleBuild.get_by_id(db.session, 8)
         assert len(module.buildrequires) == 1
         assert module.buildrequires[0].name == "platform"
         assert module.buildrequires[0].stream == "f28"
@@ -2082,7 +2082,7 @@ class TestViews:
         module_build_service.utils.load_mmd(data["modulemd"])
 
         # Make sure the buildrequires entry was created
-        module = ModuleBuild.query.get(8)
+        module = ModuleBuild.get_by_id(db.session, 8)
         assert len(module.buildrequires) == 1
         assert module.buildrequires[0].name == "platform"
         assert module.buildrequires[0].stream == "f28"
@@ -2178,7 +2178,7 @@ class TestViews:
         module_build_service.utils.load_mmd(data["modulemd"])
 
         # Make sure the buildrequires entry was created
-        module = ModuleBuild.query.get(8)
+        module = ModuleBuild.get_by_id(db.session, 8)
         assert len(module.buildrequires) == 1
         assert module.buildrequires[0].name == "platform"
         assert module.buildrequires[0].stream == "f28"