Coverage for tests/test_apdbCassandra.py: 53%

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 

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 

97class ApdbCassandraPerMonthTestCase(ApdbCassandraTestCase): 

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

99 

100 time_partition_tables = True 

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

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

103 meta_row_count = 4 

104 

105 def test_store_partition_range(self) -> None: 

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

107 config = self.make_instance() 

108 apdb = Apdb.from_config(config) 

109 

110 region = self.make_region() 

111 

112 # Visit time is beyond time_partition_end. 

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

114 catalog = makeObjectCatalog(region, 100, visit_time) 

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

116 apdb.store(visit_time, catalog) 

117 

118 # Writing to last partition makes a warning. 

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

120 catalog = makeObjectCatalog(region, 100, visit_time) 

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

122 apdb.store(visit_time, catalog) 

123 

124 

125class ApdbCassandraTestCaseReplica(ApdbCassandraTestCase): 

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

127 

128 enable_replica = True 

129 meta_row_count = 4 

130 

131 

132class ApdbCassandraTestCaseDatetimeReplica(ApdbCassandraTestCaseReplica): 

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

134 

135 use_mjd = False 

136 

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

138 if "schema_file" in kwargs: 

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

140 else: 

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

142 

143 

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

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

146 

147 time_partition_tables = False 

148 

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

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

151 kw = { 

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

153 "keyspace": self.keyspace, 

154 "schema_file": TEST_SCHEMA, 

155 "time_partition_tables": self.time_partition_tables, 

156 } 

157 kw.update(kwargs) 

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

159 

160 

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

162 """Run file leak tests.""" 

163 

164 

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

166 """Configure pytest.""" 

167 lsst.utils.tests.init() 

168 

169 

170if __name__ == "__main__": 

171 lsst.utils.tests.init() 

172 unittest.main()