Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from datetime import timedelta, datetime 

2from unittest.mock import Mock 

3 

4import pandas as pd 

5from paramiko import SFTPAttributes 

6from testfixtures import log_capture 

7 

8from sdc_etl_libs.sdc_exchange_endpoints.database_endpoints.SDCSnowflakeEndpoint import SDCSnowflakeEndpoint 

9from sdc_etl_libs.sdc_exchange_endpoints.file_endpoints.SDCSFTPEndpoint import SDCSFTPEndpoint 

10from sdc_etl_libs.sdc_filetransfer.utils.delete_files_mechanism.AfterNDaysDeleteCriteria import AfterNDaysDeleteCriteria 

11from sdc_etl_libs.sdc_filetransfer.utils.delete_files_mechanism.ExistsInTableDeleteCriteria import \ 

12 ExistsInTableDeleteCriteria 

13 

14 

15def get_after_n_days_criteria(days_threshold): 

16 return AfterNDaysDeleteCriteria(days_threshold) 

17 

18 

19def get_exists_in_table_criteria(mocker): 

20 endpoint = SDCSnowflakeEndpoint() 

21 endpoint.query = "query" 

22 endpoint.database = "database" 

23 

24 # snowflake files 

25 data = ['first_file', 'second_file'] 

26 df = pd.DataFrame(data, columns=['filename']) 

27 mocker.patch('pandas.read_sql', return_value=df) 

28 

29 return ExistsInTableDeleteCriteria(endpoint) 

30 

31 

32def get_sftp_files(): 

33 first_ftp_file = SFTPAttributes() 

34 first_ftp_file.st_mtime = (datetime.today() - timedelta(days=5)).timestamp() 

35 first_ftp_file.filename = "first_file" 

36 second_ftp_file = SFTPAttributes() 

37 second_ftp_file.st_mtime = (datetime.today() - timedelta(days=6)).timestamp() 

38 second_ftp_file.filename = "second_file" 

39 return [first_ftp_file, second_ftp_file] 

40 

41 

42@log_capture() 

43def test_delete_data_all_files(capture, mocker): 

44 # delete criteria 

45 delete_criteria_1 = get_after_n_days_criteria(3) 

46 delete_criteria_2 = get_exists_in_table_criteria(mocker) 

47 

48 # initial files 

49 initial_files = get_sftp_files() 

50 

51 # sftp endpoint config 

52 sftp_client = Mock() 

53 sftp_client.listdir_attr.return_value = initial_files 

54 sftp_endpoint = SDCSFTPEndpoint() 

55 sftp_endpoint.client = sftp_client 

56 sftp_endpoint.path = "path" 

57 

58 # test 

59 sftp_endpoint.delete_data([delete_criteria_1, delete_criteria_2]) 

60 capture.check_present( 

61 ('root', 'INFO', 'Number of files to delete from FTP Server: 2'), 

62 ('root', 'INFO', 'File first_file deleted'), 

63 ('root', 'INFO', 'File second_file deleted') 

64 ) 

65 

66 

67@log_capture() 

68def test_delete_data_one_files(capture, mocker): 

69 # delete criteria 

70 delete_criteria_1 = get_after_n_days_criteria(6) 

71 delete_criteria_2 = get_exists_in_table_criteria(mocker) 

72 

73 # initial files 

74 initial_files = get_sftp_files() 

75 

76 # sftp endpoint config 

77 sftp_client = Mock() 

78 sftp_client.listdir_attr.return_value = initial_files 

79 sftp_endpoint = SDCSFTPEndpoint() 

80 sftp_endpoint.client = sftp_client 

81 sftp_endpoint.path = "path" 

82 

83 # test 

84 sftp_endpoint.delete_data([delete_criteria_1, delete_criteria_2]) 

85 capture.check_present( 

86 ('root', 'INFO', 'Number of files to delete from FTP Server: 1'), 

87 ('root', 'INFO', 'File second_file deleted') 

88 ) 

89 

90 

91@log_capture() 

92def test_delete_data_no_files(capture, mocker): 

93 # delete criteria 

94 delete_criteria_1 = get_after_n_days_criteria(7) 

95 delete_criteria_2 = get_exists_in_table_criteria(mocker) 

96 

97 # initial files 

98 initial_files = get_sftp_files() 

99 

100 # sftp endpoint config 

101 sftp_client = Mock() 

102 sftp_client.listdir_attr.return_value = initial_files 

103 sftp_endpoint = SDCSFTPEndpoint() 

104 sftp_endpoint.client = sftp_client 

105 sftp_endpoint.path = "path" 

106 

107 # test 

108 sftp_endpoint.delete_data([delete_criteria_1, delete_criteria_2]) 

109 capture.check_present( 

110 ('root', 'INFO', 'Number of files to delete from FTP Server: 0') 

111 ) 

112 

113 

114@log_capture() 

115def test_delete_data_deletion_exception(capture, mocker): 

116 # delete criteria 

117 delete_criteria_1 = get_after_n_days_criteria(3) 

118 delete_criteria_2 = get_exists_in_table_criteria(mocker) 

119 

120 # initial files 

121 initial_files = get_sftp_files() 

122 

123 # sftp endpoint config 

124 sftp_client = Mock() 

125 sftp_client.listdir_attr.return_value = initial_files 

126 sftp_client.remove.side_effect = [Exception("delete error"), ""] 

127 sftp_endpoint = SDCSFTPEndpoint() 

128 sftp_endpoint.client = sftp_client 

129 sftp_endpoint.path = "path" 

130 

131 # test 

132 sftp_endpoint.delete_data([delete_criteria_1, delete_criteria_2]) 

133 capture.check_present( 

134 ('root', 'INFO', 'Number of files to delete from FTP Server: 2'), 

135 ('root', 'ERROR', 'Error deleting file first_file: delete error'), 

136 ('root', 'INFO', 'File second_file deleted') 

137 )