Coverage for biobb_common/biobb_common/tools/test_fixtures.py: 46%

1"""Boiler plate functions for testsys 

2""" 

3import os 

4import typing 

5from typing import Optional 

6from pathlib import Path 

7import sys 

8import shutil 

9import hashlib 

10import Bio.PDB 

11import codecs 

12from biobb_common.configuration import settings 

13from biobb_common.tools import file_utils as fu 

14import numpy as np 

15 

16 

17def test_setup(test_object, dict_key: Optional[str] = None, config: Optional[str] = None): 

18 """Add the unitest_dir, test_dir, conf_file_path, system, properties and path as 

19 attributes to the **test_object** and create a directory to launch the unitest. 

20 

21 Args: 

22 test_object (:obj:`test`): The test object. 

23 dict_key (str): Key of the test parameters in the yaml config file. 

24 config (str): Path to the configuration file. 

25 """ 

26 test_object.testfile_dir = str(Path(Path(sys.modules[test_object.__module__].__file__).resolve()).parent) 

27 test_object.unitest_dir = str(Path(test_object.testfile_dir).parent) 

28 test_object.test_dir = str(Path(test_object.unitest_dir).parent) 

29 test_object.data_dir = str(Path(test_object.test_dir).joinpath('data')) 

30 test_object.reference_dir = str(Path(test_object.test_dir).joinpath('reference')) 

31 if config: 

32 test_object.conf_file_path = config 

33 else: 

34 test_object.conf_file_path = str(Path(test_object.test_dir).joinpath('conf.yml')) 

35 

36 test_object.system = os.getenv('testsys') 

37 conf = settings.ConfReader(test_object.conf_file_path, test_object.system) 

38 

39 if dict_key: 

40 test_object.properties = conf.get_prop_dic()[dict_key] 

41 test_object.paths = {k: v.replace('test_data_dir', test_object.data_dir, 1).replace('test_reference_dir', test_object.reference_dir, 1) for k, v in conf.get_paths_dic()[dict_key].items()} 

42 else: 

43 test_object.properties = conf.get_prop_dic() 

44 test_object.paths = {k: v.replace('test_data_dir', test_object.data_dir, 1).replace('test_reference_dir', test_object.reference_dir, 1) for k, v in conf.get_paths_dic().items()} 

45 

46 fu.create_dir(test_object.properties['path']) 

47 os.chdir(test_object.properties['path']) 

48 

49 

50def test_teardown(test_object): 

51 """Remove the **test_object.properties['working_dir_path']** 

52 

53 Args: 

54 test_object (:obj:`test`): The test object. 

55 """ 

56 unitests_path = Path(test_object.properties['path']).resolve().parent 

57 print(f"\nRemoving: {unitests_path}") 

58 shutil.rmtree(unitests_path) 

59 

60 

61def exe_success(return_code: int) -> bool: 

62 """Check if **return_code** is 0 

63 

64 Args: 

65 return_code (int): Return code of a process. 

66 

67 Returns: 

68 bool: True if return code is equal to 0 

69 """ 

70 return return_code == 0 

71 

72 

73def not_empty(file_path: str) -> bool: 

74 """Check if file exists and is not empty. 

75 

76 Args: 

77 file_path (str): Path to the file. 

78 

79 Returns: 

80 bool: True if **file_path** exists and is not empty. 

81 """ 

82 print("Checking if empty file: "+file_path) 

83 return Path(file_path).is_file() and Path(file_path).stat().st_size > 0 

84 

85 

86def compare_hash(file_a: str, file_b: str) -> bool: 

87 """Compute and compare the hashes of two files""" 

88 print("Comparing: ") 

89 print(" File_A: "+file_a) 

90 print(" File_B: "+file_b) 

91 file_a_hash = hashlib.sha256(open(file_a, 'rb').read()).digest() 

92 file_b_hash = hashlib.sha256(open(file_b, 'rb').read()).digest() 

93 print(" File_A hash: "+str(file_a_hash)) 

94 print(" File_B hash: "+str(file_b_hash)) 

95 return file_a_hash == file_b_hash 

96 

97 

98def equal(file_a: str, file_b: str, ignore_list: Optional[typing.List[typing.Union[str, int]]] = None, **kwargs) -> bool: 

99 """Check if two files are equal""" 

100 if ignore_list: 

101 # Line by line comparison 

102 return compare_line_by_line(file_a, file_b, ignore_list) 

103 

104 if file_a.endswith(".zip") and file_b.endswith(".zip"): 

105 return compare_zip(file_a, file_b) 

106 

107 if file_a.endswith(".pdb") and file_b.endswith(".pdb"): 

108 return compare_pdb(file_a, file_b, **kwargs) 

109 

110 if file_a.endswith(".top") and file_b.endswith(".top"): 

111 return compare_top_itp(file_a, file_b) 

112 

113 if file_a.endswith(".itp") and file_b.endswith(".itp"): 

114 return compare_top_itp(file_a, file_b) 

115 

116 if file_a.endswith(".gro") and file_b.endswith(".gro"): 

117 return compare_ignore_first(file_a, file_b) 

118 

119 if file_a.endswith(".prmtop") and file_b.endswith(".prmtop"): 

120 return compare_ignore_first(file_a, file_b) 

121 

122 if file_a.endswith(".inp") and file_b.endswith(".inp"): 

123 return compare_ignore_first(file_a, file_b) 

124 

125 if file_a.endswith(".par") and file_b.endswith(".par"): 

126 return compare_ignore_first(file_a, file_b) 

127 

128 if file_a.endswith((".nc", ".netcdf", ".xtc")) and file_b.endswith((".nc", ".netcdf", ".xtc")): 

129 return compare_size(file_a, file_b, kwargs.get('percent_tolerance', 1.0)) 

130 

131 if file_a.endswith(".xvg") and file_b.endswith(".xvg"): 

132 return compare_xvg(file_a, file_b, kwargs.get('percent_tolerance', 1.0)) 

133 

134 image_extensions = ('.png', '.jfif', '.ppm', '.tiff', '.jpg', '.dib', '.pgm', '.bmp', '.jpeg', '.pbm', '.jpe', '.apng', '.pnm', '.gif', '.tif') 

135 if file_a.endswith(image_extensions) and file_b.endswith(image_extensions): 

136 return compare_images(file_a, file_b, kwargs.get('percent_tolerance', 1.0)) 

137 

138 return compare_hash(file_a, file_b) 

139 

140 

141def compare_line_by_line(file_a: str, file_b: str, ignore_list: typing.List[typing.Union[str, int]]) -> bool: 

142 with open(file_a) as fa, open(file_b) as fb: 

143 for index, (line_a, line_b) in enumerate(zip(fa, fb)): 

144 if index in ignore_list or any(word in line_a for word in ignore_list if isinstance(word, str)): 

145 continue 

146 elif line_a != line_b: 

147 return False 

148 return True 

149 

150 

151def equal_txt(file_a: str, file_b: str) -> bool: 

152 """Check if two text files are equal""" 

153 return compare_hash(file_a, file_b) 

154 

155 

156def compare_zip(zip_a: str, zip_b: str) -> bool: 

157 """ Compare zip files """ 

158 print("This is a ZIP comparison!") 

159 print("Unzipping:") 

160 print("Creating a unique_dir for: %s" % zip_a) 

161 zip_a_dir = fu.create_unique_dir() 

162 zip_a_list = fu.unzip_list(zip_a, dest_dir=zip_a_dir) 

163 print("Creating a unique_dir for: %s" % zip_b) 

164 zip_b_dir = fu.create_unique_dir() 

165 zip_b_list = fu.unzip_list(zip_b, dest_dir=zip_b_dir) 

166 

167 if not len(zip_a_list) == len(zip_b_list): 

168 return False 

169 

170 for uncompressed_zip_a in zip_a_list: 

171 uncompressed_zip_b = str(Path(zip_b_dir).joinpath(Path(uncompressed_zip_a).name)) 

172 if not equal(uncompressed_zip_a, uncompressed_zip_b): 

173 return False 

174 

175 return True 

176 

177 

178def compare_pdb(pdb_a: str, pdb_b: str, rmsd_cutoff: int = 1, remove_hetatm: bool = True, remove_hydrogen: bool = True, **kwargs): 

179 """ Compare pdb files """ 

180 print("Checking RMSD between:") 

181 print(" PDB_A: "+pdb_a) 

182 print(" PDB_B: "+pdb_b) 

183 pdb_parser = Bio.PDB.PDBParser(PERMISSIVE=True, QUIET=True) 

184 st_a = pdb_parser.get_structure("st_a", pdb_a)[0] 

185 st_b = pdb_parser.get_structure("st_b", pdb_b)[0] 

186 

187 if remove_hetatm: 

188 print(" Ignoring HETAMT in RMSD") 

189 residues_a = [list(res.get_atoms()) for res in st_a.get_residues() if not res.id[0].startswith('H_')] 

190 residues_b = [list(res.get_atoms()) for res in st_b.get_residues() if not res.id[0].startswith('H_')] 

191 atoms_a = [atom for residue in residues_a for atom in residue] 

192 atoms_b = [atom for residue in residues_b for atom in residue] 

193 else: 

194 atoms_a = st_a.get_atoms() 

195 atoms_b = st_b.get_atoms() 

196 

197 if remove_hydrogen: 

198 print(" Ignoring Hydrogen atoms in RMSD") 

199 atoms_a = [atom for atom in atoms_a if not atom.get_name().startswith('H')] 

200 atoms_b = [atom for atom in atoms_b if not atom.get_name().startswith('H')] 

201 

202 print(" Atoms ALIGNED in PDB_A: "+str(len(atoms_a))) 

203 print(" Atoms ALIGNED in PDB_B: "+str(len(atoms_b))) 

204 super_imposer = Bio.PDB.Superimposer() 

205 super_imposer.set_atoms(atoms_a, atoms_b) 

206 super_imposer.apply(atoms_b) 

207 print(' RMS: '+str(super_imposer.rms)) 

208 print(' RMS_CUTOFF: '+str(rmsd_cutoff)) 

209 return super_imposer.rms < rmsd_cutoff 

210 

211 

212def compare_top_itp(file_a: str, file_b: str) -> bool: 

213 """ Compare top/itp files """ 

214 print("Comparing TOP/ITP:") 

215 print(" FILE_A: "+file_a) 

216 print(" FILE_B: "+file_b) 

217 with codecs.open(file_a, 'r', encoding='utf-8', errors='ignore') as f_a: 

218 next(f_a) 

219 with codecs.open(file_b, 'r', encoding='utf-8', errors='ignore') as f_b: 

220 next(f_b) 

221 return [line.strip() for line in f_a if not line.strip().startswith(';')] == [line.strip() for line in f_b if not line.strip().startswith(';')] 

222 

223 

224def compare_ignore_first(file_a: str, file_b: str) -> bool: 

225 """ Compare two files ignoring the first line """ 

226 print("Comparing ignoring first line of both files:") 

227 print(" FILE_A: "+file_a) 

228 print(" FILE_B: "+file_b) 

229 with open(file_a) as f_a: 

230 next(f_a) 

231 with open(file_b) as f_b: 

232 next(f_b) 

233 return [line.strip() for line in f_a] == [line.strip() for line in f_b] 

234 

235 

236def compare_size(file_a: str, file_b: str, percent_tolerance: float = 1.0) -> bool: 

237 """ Compare two files using size """ 

238 print("Comparing size of both files:") 

239 print(f" FILE_A: {file_a}") 

240 print(f" FILE_B: {file_b}") 

241 size_a = Path(file_a).stat().st_size 

242 size_b = Path(file_b).stat().st_size 

243 average_size = (size_a + size_b) / 2 

244 tolerance = average_size * percent_tolerance / 100 

245 tolerance_low = average_size - tolerance 

246 tolerance_high = average_size + tolerance 

247 print(f" SIZE_A: {size_a} bytes") 

248 print(f" SIZE_B: {size_b} bytes") 

249 print(f" TOLERANCE: {percent_tolerance}%, Low: {tolerance_low} bytes, High: {tolerance_high} bytes") 

250 return (tolerance_low <= size_a <= tolerance_high) and (tolerance_low <= size_b <= tolerance_high) 

251 

252 

253def compare_xvg(file_a: str, file_b: str, percent_tolerance: float = 1.0) -> bool: 

254 """ Compare two files using size """ 

255 print("Comparing size of both files:") 

256 print(f" FILE_A: {file_a}") 

257 print(f" FILE_B: {file_b}") 

258 arrays_tuple_a = np.loadtxt(file_a, comments="@", unpack=True) 

259 arrays_tuple_b = np.loadtxt(file_b, comments="@", unpack=True) 

260 for array_a, array_b in zip(arrays_tuple_a, arrays_tuple_b): 

261 if not np.allclose(array_a, array_b, rtol=percent_tolerance / 100): 

262 return False 

263 return True 

264 

265 

266def compare_images(file_a: str, file_b: str, percent_tolerance: float = 1.0) -> bool: 

267 try: 

268 from PIL import Image 

269 import imagehash 

270 except ImportError: 

271 print("To compare images, please install the following packages: Pillow, imagehash") 

272 return False 

273 

274 """ Compare two files using size """ 

275 print("Comparing images of both files:") 

276 print(f" IMAGE_A: {file_a}") 

277 print(f" IMAGE_B: {file_b}") 

278 hash_a = imagehash.average_hash(Image.open(file_a)) 

279 hash_b = imagehash.average_hash(Image.open(file_b)) 

280 tolerance = (len(hash_a) + len(hash_b)) / 2 * percent_tolerance / 100 

281 if tolerance < 1: 

282 tolerance = 1 

283 difference = hash_a - hash_b 

284 print(f" IMAGE_A HASH: {hash_a} SIZE: {len(hash_a)} bits") 

285 print(f" IMAGE_B HASH: {hash_b} SIZE: {len(hash_b)} bits") 

286 print(f" TOLERANCE: {percent_tolerance}%, ABS TOLERANCE: {tolerance} bits, DIFFERENCE: {difference} bits") 

287 if difference > tolerance: 

288 return False 

289 return True