Coverage for tests/test_butler_query.py: 52%

74 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-16 10:44 +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/>. 

27 

28from __future__ import annotations 

29 

30import gc 

31import os 

32import unittest 

33from typing import Any 

34 

35import sqlalchemy 

36 

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 

45 

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 

50 

51TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

52 

53 

54class ButlerQueryTestsDirectSQLite(ButlerQueryTests, unittest.TestCase): 

55 """Unit tests for DirectButler.query methods with sqlite database.""" 

56 

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 

61 

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() 

68 

69 def tearDown(self) -> None: 

70 removeTestTempDir(self.root) 

71 super().tearDown() 

72 

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) 

77 

78 for filename in args: 

79 self.load_data(butler._registry, filename) 

80 

81 self.make_bias_collection(butler._registry) 

82 

83 return butler 

84 

85 

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.""" 

89 

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 

94 

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;")) 

100 

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() 

108 

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() 

116 

117 def setUp(self) -> None: 

118 self.server = self.postgresql() 

119 

120 self.root = makeTestTempDir(TESTDIR) 

121 self.tmpConfigFile = os.path.join(self.root, "butler.yaml") 

122 

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) 

127 

128 super().setUp() 

129 

130 def tearDown(self) -> None: 

131 self.server.stop() 

132 removeTestTempDir(self.root) 

133 super().tearDown() 

134 

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) 

139 

140 for filename in args: 

141 self.load_data(butler._registry, filename) 

142 

143 self.make_bias_collection(butler._registry) 

144 

145 return butler 

146 

147 

148if __name__ == "__main__": 

149 unittest.main()