Coverage for tests/test_apdbCassandra.py: 49%

74 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-06 04:04 -0700

1# This file is part of dax_apdb. 

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 program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22"""Unit test for `ApdbCassandra` class. 

23 

24Notes 

25----- 

26For now this test can only run against actual Cassandra cluster, to specify 

27cluster location use ``DAX_APDB_TEST_CASSANDRA_CLUSTER`` environment variable, 

28e.g.: 

29 

30 export DAX_APDB_TEST_CASSANDRA_CLUSTER=cassandra.example.com 

31 pytest tests/test_apdbCassandra.py 

32 

33Individual tests create and destroy unique keyspaces in the cluster, there is 

34no need to pre-create a keyspace with predefined name. 

35""" 

36 

37import logging 

38import os 

39import unittest 

40import uuid 

41from typing import TYPE_CHECKING, Any 

42 

43try: 

44 from cassandra.cluster import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile 

45 from cassandra.policies import RoundRobinPolicy 

46 

47 CASSANDRA_IMPORTED = True 

48except ImportError: 

49 CASSANDRA_IMPORTED = False 

50 

51import lsst.utils.tests 

52from lsst.dax.apdb import ApdbCassandra, ApdbCassandraConfig, ApdbConfig, ApdbTables 

53from lsst.dax.apdb.apdbCassandra import CASSANDRA_IMPORTED 

54from lsst.dax.apdb.tests import ApdbSchemaUpdateTest, ApdbTest 

55 

56TEST_SCHEMA = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema.yaml") 

57 

58logging.basicConfig(level=logging.INFO) 

59 

60 

61class ApdbCassandraMixin: 

62 """Mixin class which defines common methods for unit tests.""" 

63 

64 schema_path = TEST_SCHEMA 

65 

66 @classmethod 

67 def setUpClass(cls) -> None: 

68 """Prepare config for server connection.""" 

69 if not CASSANDRA_IMPORTED: 

70 raise unittest.SkipTest("FAiled to import Cassandra modules") 

71 cluster_host = os.environ.get("DAX_APDB_TEST_CASSANDRA_CLUSTER") 

72 if not cluster_host: 

73 raise unittest.SkipTest("DAX_APDB_TEST_CASSANDRA_CLUSTER is not set") 

74 if not CASSANDRA_IMPORTED: 

75 raise unittest.SkipTest("cassandra_driver cannot be imported") 

76 

77 def _run_query(self, query: str) -> None: 

78 # Used protocol version from default config. 

79 config = ApdbCassandraConfig() 

80 default_profile = ExecutionProfile(load_balancing_policy=RoundRobinPolicy()) 

81 profiles = {EXEC_PROFILE_DEFAULT: default_profile} 

82 cluster = Cluster( 

83 contact_points=[self.cluster_host], 

84 execution_profiles=profiles, 

85 protocol_version=config.protocol_version, 

86 ) 

87 session = cluster.connect() 

88 session.execute(query, timeout=120) 

89 del session 

90 cluster.shutdown() 

91 

92 def setUp(self) -> None: 

93 """Prepare config for server connection.""" 

94 self.cluster_host = os.environ.get("DAX_APDB_TEST_CASSANDRA_CLUSTER") 

95 # Use dedicated keyspace for each test, keyspace is created by 

96 # init_database if it does not exist. 

97 key = uuid.uuid4() 

98 self.keyspace = f"apdb_{key.hex}" 

99 

100 def tearDown(self) -> None: 

101 # Delete per-test keyspace. 

102 query = f"DROP KEYSPACE {self.keyspace}" 

103 self._run_query(query) 

104 

105 if TYPE_CHECKING: 

106 # For mypy. 

107 def make_instance(self, **kwargs: Any) -> ApdbConfig: ... 

108 

109 

110class ApdbCassandraTestCase(ApdbCassandraMixin, ApdbTest, unittest.TestCase): 

111 """A test case for ApdbCassandra class""" 

112 

113 allow_visit_query = False 

114 time_partition_tables = False 

115 time_partition_start: str | None = None 

116 time_partition_end: str | None = None 

117 

118 def make_instance(self, **kwargs: Any) -> ApdbConfig: 

119 """Make config class instance used in all tests.""" 

120 kw: dict[str, Any] = { 

121 "hosts": [self.cluster_host], 

122 "keyspace": self.keyspace, 

123 "schema_file": TEST_SCHEMA, 

124 "time_partition_tables": self.time_partition_tables, 

125 "use_insert_id": self.use_insert_id, 

126 } 

127 if self.time_partition_start: 

128 kw["time_partition_start"] = self.time_partition_start 

129 if self.time_partition_end: 

130 kw["time_partition_end"] = self.time_partition_end 

131 kw.update(kwargs) 

132 return ApdbCassandra.init_database(**kw) 

133 

134 def getDiaObjects_table(self) -> ApdbTables: 

135 """Return type of table returned from getDiaObjects method.""" 

136 return ApdbTables.DiaObjectLast 

137 

138 

139class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

140 """A test case for ApdbCassandra class with per-month tables.""" 

141 

142 time_partition_tables = True 

143 time_partition_start = "2019-12-01T00:00:00" 

144 time_partition_end = "2022-01-01T00:00:00" 

145 

146 

147class ApdbCassandraTestCaseInsertIds(ApdbCassandraTestCase): 

148 """A test case with use_insert_id.""" 

149 

150 use_insert_id = True 

151 

152 

153class ApdbSchemaUpdateCassandraTestCase(ApdbCassandraMixin, ApdbSchemaUpdateTest, unittest.TestCase): 

154 """A test case for schema updates using Cassandra backend.""" 

155 

156 def make_instance(self, **kwargs: Any) -> ApdbConfig: 

157 """Make config class instance used in all tests.""" 

158 kw = { 

159 "hosts": [self.cluster_host], 

160 "keyspace": self.keyspace, 

161 "schema_file": TEST_SCHEMA, 

162 "time_partition_tables": False, 

163 } 

164 kw.update(kwargs) 

165 return ApdbCassandra.init_database(**kw) # type: ignore[arg-type] 

166 

167 

168class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

169 """Run file leak tests.""" 

170 

171 

172def setup_module(module: Any) -> None: 

173 """Configure pytest.""" 

174 lsst.utils.tests.init() 

175 

176 

177if __name__ == "__main__": 

178 lsst.utils.tests.init() 

179 unittest.main()