Coverage for tests/utils.py: 92%

1import os 

2import warnings 

3 

4import peprs 

5import yaml 

6from sqlalchemy.exc import OperationalError 

7 

8from pepdbagent import PEPDatabaseAgent 

9 

10DSN = "postgresql+psycopg://postgres:pass8743hf9h23f87h437@localhost:5432/pep-db" 

11 

12TESTS_PATH = os.path.join( 

13 os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 

14 "tests", 

15) 

16 

17DATA_PATH = os.path.join( 

18 TESTS_PATH, 

19 "data", 

20) 

21 

22SCHEMAS_PATH = os.path.join( 

23 TESTS_PATH, 

24 "schemas", 

25) 

26 

27 

28def get_path_to_example_file(namespace: str, project_name: str) -> str: 

29 """ 

30 Get path to example file 

31 """ 

32 return os.path.join(DATA_PATH, namespace, project_name, "project_config.yaml") 

33 

34 

35def get_path_to_example_schema(namespace: str, schema_name: str) -> str: 

36 """ 

37 Get path to example schema 

38 """ 

39 return os.path.join(SCHEMAS_PATH, namespace, schema_name) 

40 

41 

42def list_of_available_peps() -> dict: 

43 pep_namespaces = os.listdir(DATA_PATH) 

44 projects = {} 

45 for np in pep_namespaces: 

46 pep_name = os.listdir(os.path.join(DATA_PATH, np)) 

47 projects[np] = {p: get_path_to_example_file(np, p) for p in pep_name} 

48 return projects 

49 

50 

51def list_of_available_schemas() -> dict: 

52 schema_namespaces = os.listdir(SCHEMAS_PATH) 

53 schemas = {} 

54 for np in schema_namespaces: 

55 schema_name = os.listdir(os.path.join(SCHEMAS_PATH, np)) 

56 schemas[np] = {p: get_path_to_example_schema(np, p) for p in schema_name} 

57 return schemas 

58 

59 

60def read_yaml_file(file_path: str) -> dict: 

61 """ 

62 Read yaml file 

63 """ 

64 with open(file_path, "r") as file: 

65 return yaml.safe_load(file) 

66 

67 

68class PEPDBAgentContextManager: 

69 """ 

70 Class with context manager to connect to database. Adds data and drops everything from the database upon exit to ensure. 

71 """ 

72 

73 def __init__( 

74 self, url: str = DSN, add_data: bool = False, add_schemas=True, echo=False 

75 ): 

76 """ 

77 :param url: database url e.g. "postgresql+psycopg://postgres:docker@localhost:5432/pep-db" 

78 :param add_data: add data to the database 

79 """ 

80 

81 self.url = url 

82 self._agent = None 

83 self._echo = echo 

84 self.add_data = add_data 

85 self.add_schemas = add_schemas 

86 

87 def __enter__(self): 

88 self._agent = PEPDatabaseAgent(dsn=self.url, echo=False) 

89 self.db_engine = self._agent.pep_db_engine 

90 self.db_engine.create_schema() 

91 if self.add_schemas: 

92 self._add_schemas() 

93 if self.add_data: 

94 self._insert_data() 

95 return self._agent 

96 

97 def __exit__(self, exc_type, exc_value, exc_traceback): 

98 self.db_engine.delete_schema() 

99 

100 def __del__(self): 

101 self.db_engine.delete_schema() 

102 

103 def _insert_data(self): 

104 pepdb_con = PEPDatabaseAgent(dsn=self.url, echo=self._echo) 

105 for namespace, item in list_of_available_peps().items(): 

106 if namespace == "private_test": 

107 private = True 

108 else: 

109 private = False 

110 for name, path in item.items(): 

111 try: 

112 prj = peprs.Project(path) 

113 except Exception as e: 

114 raise RuntimeError( 

115 f"Failed to load test project {namespace}/{name} from {path}" 

116 ) from e 

117 pepdb_con.project.create( 

118 namespace=namespace, 

119 name=name, 

120 tag="default", 

121 is_private=private, 

122 project=prj, 

123 overwrite=True, 

124 pep_schema="namespace1/2.0.0:1.0.0", # TODO: test without this line 

125 ) 

126 

127 def _add_schemas(self): 

128 pepdb_con = PEPDatabaseAgent(dsn=self.url, echo=self._echo) 

129 for namespace, item in list_of_available_schemas().items(): 

130 for name, path in item.items(): 

131 file_dict = read_yaml_file(path) 

132 

133 pepdb_con.schema.create( 

134 namespace=namespace, 

135 name=name[0:-5], 

136 version="1.0.0", 

137 schema_value=file_dict, 

138 maintainers="Teddy", 

139 contributors="Teddy, John", 

140 release_notes="Initial release", 

141 tags={"maturity_level": "trial_use"}, 

142 ) 

143 

144 @property 

145 def agent(self) -> PEPDatabaseAgent: 

146 return self._agent 

147 

148 def db_setup(self): 

149 # Check if the database is setup 

150 try: 

151 PEPDatabaseAgent(dsn=self.url) 

152 except OperationalError: 

153 warnings.warn( 

154 UserWarning( 

155 f"Skipping tests, because DB is not setup. {self.url}. To setup DB go to README.md" 

156 ) 

157 ) 

158 return False 

159 return True