Coverage for tests/test_apdbCassandra.py: 51%

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 os 

38import unittest 

39from typing import Any 

40 

41import astropy.time 

42 

43import lsst.utils.tests 

44from lsst.dax.apdb import Apdb, ApdbConfig, ApdbTables, ApdbUpdateRecord, ReplicaChunk 

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

46from lsst.dax.apdb.pixelization import Pixelization 

47from lsst.dax.apdb.tests import ApdbSchemaUpdateTest, ApdbTest, cassandra_mixin 

48from lsst.dax.apdb.tests.data_factory import makeObjectCatalog 

49 

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

51# Schema that uses `datetime` for timestamps. 

52TEST_SCHEMA_DT = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config/schema-datetime.yaml") 

53 

54 

55class ApdbCassandraMixin(cassandra_mixin.ApdbCassandraMixin): 

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

57 

58 def pixelization(self, config: ApdbConfig) -> Pixelization: 

59 """Return pixelization used by implementation.""" 

60 assert isinstance(config, ApdbCassandraConfig), "Only expect ApdbCassandraConfig here" 

61 return Pixelization( 

62 config.partitioning.part_pixelization, 

63 config.partitioning.part_pix_level, 

64 config.partitioning.part_pix_max_ranges, 

65 ) 

66 

67 

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

69 """A test case for ApdbCassandra class""" 

70 

71 time_partition_tables = False 

72 time_partition_start: str | None = None 

73 time_partition_end: str | None = None 

74 extra_chunk_columns = 2 

75 

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

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

78 kw: dict[str, Any] = { 

79 "hosts": (self.cluster_host,), 

80 "keyspace": self.keyspace, 

81 "schema_file": TEST_SCHEMA, 

82 "time_partition_tables": self.time_partition_tables, 

83 "enable_replica": self.enable_replica, 

84 } 

85 if self.time_partition_start: 

86 kw["time_partition_start"] = self.time_partition_start 

87 if self.time_partition_end: 

88 kw["time_partition_end"] = self.time_partition_end 

89 kw.update(kwargs) 

90 return ApdbCassandra.init_database(**kw) 

91 

92 def getDiaObjects_table(self) -> ApdbTables: 

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

94 return ApdbTables.DiaObjectLast 

95 

96 def store_update_records(self, apdb: Apdb, records: list[ApdbUpdateRecord], chunk: ReplicaChunk) -> None: 

97 # Docstring inherited. 

98 assert isinstance(apdb, ApdbCassandra), "Expecting ApdbCassandra instance" 

99 apdb._storeUpdateRecords(records, chunk, store_chunk=True) 

100 

101 

102class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

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

104 

105 time_partition_tables = True 

106 time_partition_start = "2020-06-01T00:00:00" 

107 time_partition_end = "2021-06-01T00:00:00" 

108 meta_row_count = 4 

109 

110 def test_store_partition_range(self) -> None: 

111 """Test that writing to non-existing partition raises an error.""" 

112 config = self.make_instance() 

113 apdb = Apdb.from_config(config) 

114 

115 region = self.make_region() 

116 

117 # Visit time is beyond time_partition_end. 

118 visit_time = astropy.time.Time("2022-01-01", format="isot", scale="tai") 

119 catalog = makeObjectCatalog(region, 100, visit_time) 

120 with self.assertRaisesRegex(ValueError, "time partitions that do not yet exist"): 

121 apdb.store(visit_time, catalog) 

122 

123 # Writing to last partition makes a warning. 

124 visit_time = astropy.time.Time("2021-06-01", format="isot", scale="tai") 

125 catalog = makeObjectCatalog(region, 100, visit_time) 

126 with self.assertWarnsRegex(UserWarning, "Writing into the last temporal partition"): 

127 apdb.store(visit_time, catalog) 

128 

129 

130class ApdbCassandraTestCaseReplica(ApdbCassandraTestCase): 

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

132 

133 enable_replica = True 

134 meta_row_count = 4 

135 

136 

137class ApdbCassandraTestCaseDatetimeReplica(ApdbCassandraTestCaseReplica): 

138 """A test case with datetime timestamps.""" 

139 

140 use_mjd = False 

141 

142 def setUp(self) -> None: 

143 super().setUp() 

144 # Schema for datetime case is also missing a validityTime column in 

145 # DiaObjectLast table. 

146 self.table_column_count[ApdbTables.DiaObjectLast] = 4 

147 

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

149 if "schema_file" in kwargs: 

150 return super().make_instance(**kwargs) 

151 else: 

152 return super().make_instance(schema_file=TEST_SCHEMA_DT, **kwargs) 

153 

154 

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

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

157 

158 time_partition_tables = False 

159 

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

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

162 kw = { 

163 "hosts": (self.cluster_host,), 

164 "keyspace": self.keyspace, 

165 "schema_file": TEST_SCHEMA, 

166 "time_partition_tables": self.time_partition_tables, 

167 } 

168 kw.update(kwargs) 

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

170 

171 

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

173 """Run file leak tests.""" 

174 

175 

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

177 """Configure pytest.""" 

178 lsst.utils.tests.init() 

179 

180 

181if __name__ == "__main__": 

182 lsst.utils.tests.init() 

183 unittest.main()