Coverage for tests/test_apdbCassandra.py: 48%

73 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 03:20 -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 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 

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

108 """A test case for ApdbCassandra class""" 

109 

110 time_partition_tables = False 

111 time_partition_start: str | None = None 

112 time_partition_end: str | None = None 

113 

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

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

116 kw: dict[str, Any] = { 

117 "hosts": [self.cluster_host], 

118 "keyspace": self.keyspace, 

119 "schema_file": TEST_SCHEMA, 

120 "time_partition_tables": self.time_partition_tables, 

121 "use_insert_id": self.enable_replica, 

122 } 

123 if self.time_partition_start: 

124 kw["time_partition_start"] = self.time_partition_start 

125 if self.time_partition_end: 

126 kw["time_partition_end"] = self.time_partition_end 

127 kw.update(kwargs) 

128 return ApdbCassandra.init_database(**kw) 

129 

130 def getDiaObjects_table(self) -> ApdbTables: 

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

132 return ApdbTables.DiaObjectLast 

133 

134 

135class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

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

137 

138 time_partition_tables = True 

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

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

141 

142 

143class ApdbCassandraTestCaseReplica(ApdbCassandraTestCase): 

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

145 

146 enable_replica = True 

147 

148 

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

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

151 

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

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

154 kw = { 

155 "hosts": [self.cluster_host], 

156 "keyspace": self.keyspace, 

157 "schema_file": TEST_SCHEMA, 

158 "time_partition_tables": False, 

159 } 

160 kw.update(kwargs) 

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

162 

163 

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

165 """Run file leak tests.""" 

166 

167 

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

169 """Configure pytest.""" 

170 lsst.utils.tests.init() 

171 

172 

173if __name__ == "__main__": 

174 lsst.utils.tests.init() 

175 unittest.main()