Coverage for tests/test_apdbCassandra.py: 49%

74 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-20 02:41 -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 ApdbConfig, ApdbTables 

53from lsst.dax.apdb.cassandra import ApdbCassandra, ApdbCassandraConfig 

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 # Deleting many tables can take long time, use long timeout. 

89 session.execute(query, timeout=600) 

90 del session 

91 cluster.shutdown() 

92 

93 def setUp(self) -> None: 

94 """Prepare config for server connection.""" 

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

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

97 # init_database if it does not exist. 

98 key = uuid.uuid4() 

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

100 

101 def tearDown(self) -> None: 

102 # Delete per-test keyspace. 

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

104 self._run_query(query) 

105 

106 if TYPE_CHECKING: 

107 # For mypy. 

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

109 

110 

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

112 """A test case for ApdbCassandra class""" 

113 

114 allow_visit_query = False 

115 time_partition_tables = False 

116 time_partition_start: str | None = None 

117 time_partition_end: str | None = None 

118 

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

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

121 kw: dict[str, Any] = { 

122 "hosts": [self.cluster_host], 

123 "keyspace": self.keyspace, 

124 "schema_file": TEST_SCHEMA, 

125 "time_partition_tables": self.time_partition_tables, 

126 "use_insert_id": self.enable_replica, 

127 } 

128 if self.time_partition_start: 

129 kw["time_partition_start"] = self.time_partition_start 

130 if self.time_partition_end: 

131 kw["time_partition_end"] = self.time_partition_end 

132 kw.update(kwargs) 

133 return ApdbCassandra.init_database(**kw) 

134 

135 def getDiaObjects_table(self) -> ApdbTables: 

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

137 return ApdbTables.DiaObjectLast 

138 

139 

140class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

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

142 

143 time_partition_tables = True 

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

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

146 

147 

148class ApdbCassandraTestCaseReplica(ApdbCassandraTestCase): 

149 """A test case with enabled replica tables.""" 

150 

151 enable_replica = True 

152 

153 

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

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

156 

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

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

159 kw = { 

160 "hosts": [self.cluster_host], 

161 "keyspace": self.keyspace, 

162 "schema_file": TEST_SCHEMA, 

163 "time_partition_tables": False, 

164 } 

165 kw.update(kwargs) 

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

167 

168 

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

170 """Run file leak tests.""" 

171 

172 

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

174 """Configure pytest.""" 

175 lsst.utils.tests.init() 

176 

177 

178if __name__ == "__main__": 

179 lsst.utils.tests.init() 

180 unittest.main()