Coverage for tests/test_butler_query.py: 52%
74 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-13 10:57 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-13 10:57 +0000
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30import gc
31import os
32import unittest
33from typing import Any
35import sqlalchemy
37try:
38 # It's possible but silly to have testing.postgresql installed without
39 # having the postgresql server installed (because then nothing in
40 # testing.postgresql would work), so we use the presence of that module
41 # to test whether we can expect the server to be available.
42 import testing.postgresql # type: ignore[import]
43except ImportError:
44 testing = None
46from lsst.daf.butler import Butler, Config
47from lsst.daf.butler.direct_butler import DirectButler
48from lsst.daf.butler.tests.butler_query import ButlerQueryTests
49from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
51TESTDIR = os.path.abspath(os.path.dirname(__file__))
54class ButlerQueryTestsDirectSQLite(ButlerQueryTests, unittest.TestCase):
55 """Unit tests for DirectButler.query methods with sqlite database."""
57 data_dir = os.path.join(TESTDIR, "data/registry")
58 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
59 root: str
60 tmpConfigFile: str
62 def setUp(self) -> None:
63 """Create a new butler root for each test."""
64 self.root = makeTestTempDir(TESTDIR)
65 Butler.makeRepo(self.root, config=Config(self.configFile))
66 self.tmpConfigFile = os.path.join(self.root, "butler.yaml")
67 super().setUp()
69 def tearDown(self) -> None:
70 removeTestTempDir(self.root)
71 super().tearDown()
73 def make_butler(self, *args: str) -> Butler:
74 # Docstring inherited.
75 butler = Butler.from_config(self.tmpConfigFile, run="run")
76 assert isinstance(butler, DirectButler)
78 for filename in args:
79 self.load_data(butler._registry, filename)
81 self.make_bias_collection(butler._registry)
83 return butler
86@unittest.skipUnless(testing is not None, "testing.postgresql module not found")
87class ButlerQueryTestsDirectPostgres(ButlerQueryTests, unittest.TestCase):
88 """Unit tests for DirectButler.query methods with postgres database."""
90 data_dir = os.path.join(TESTDIR, "data/registry")
91 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")
92 root: str
93 tmpConfigFile: str
95 @staticmethod
96 def _handler(postgresql: Any) -> None:
97 engine = sqlalchemy.engine.create_engine(postgresql.url())
98 with engine.begin() as connection:
99 connection.execute(sqlalchemy.text("CREATE EXTENSION btree_gist;"))
101 @classmethod
102 def setUpClass(cls) -> None:
103 # Create the postgres test server.
104 cls.postgresql = testing.postgresql.PostgresqlFactory(
105 cache_initialized_db=True, on_initialized=cls._handler
106 )
107 super().setUpClass()
109 @classmethod
110 def tearDownClass(cls) -> None:
111 # Clean up any lingering SQLAlchemy engines/connections
112 # so they're closed before we shut down the server.
113 gc.collect()
114 cls.postgresql.clear_cache()
115 super().tearDownClass()
117 def setUp(self) -> None:
118 self.server = self.postgresql()
120 self.root = makeTestTempDir(TESTDIR)
121 self.tmpConfigFile = os.path.join(self.root, "butler.yaml")
123 # Use default config but update database URI.
124 config = Config(self.configFile)
125 config["registry", "db"] = self.server.url()
126 Butler.makeRepo(self.root, config=config)
128 super().setUp()
130 def tearDown(self) -> None:
131 self.server.stop()
132 removeTestTempDir(self.root)
133 super().tearDown()
135 def make_butler(self, *args: str) -> Butler:
136 # Docstring inherited.
137 butler = Butler.from_config(self.tmpConfigFile, run="run")
138 assert isinstance(butler, DirectButler)
140 for filename in args:
141 self.load_data(butler._registry, filename)
143 self.make_bias_collection(butler._registry)
145 return butler
148if __name__ == "__main__":
149 unittest.main()