Coverage for python / felis / tests / postgresql.py: 45%

1"""Provides a temporary Postgresql instance for testing.""" 

2 

3# This file is part of felis. 

4# 

5# Developed for the LSST Data Management System. 

6# This product includes software developed by the LSST Project 

7# (https://www.lsst.org). 

8# See the COPYRIGHT file at the top-level directory of this distribution 

9# for details of code ownership. 

10# 

11# This program is free software: you can redistribute it and/or modify 

12# it under the terms of the GNU General Public License as published by 

13# the Free Software Foundation, either version 3 of the License, or 

14# (at your option) any later version. 

15# 

16# This program is distributed in the hope that it will be useful, 

17# but WITHOUT ANY WARRANTY; without even the implied warranty of 

18# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

19# GNU General Public License for more details. 

20# 

21# You should have received a copy of the GNU General Public License 

22# along with this program. If not, see <https://www.gnu.org/licenses/>. 

23 

24import gc 

25import unittest 

26from collections.abc import Iterator 

27from contextlib import contextmanager 

28 

29from sqlalchemy import text 

30from sqlalchemy.engine import Connection, Engine, create_engine 

31 

32try: 

33 from testing.postgresql import Postgresql # type: ignore 

34except ImportError: 

35 Postgresql = None 

36 

37__all__ = ["TemporaryPostgresInstance", "setup_postgres_test_db"] 

38 

39 

40class TemporaryPostgresInstance: 

41 """Wrapper for a temporary Postgres database. 

42 

43 Parameters 

44 ---------- 

45 server 

46 The ``testing.postgresql.Postgresql`` instance. 

47 engine 

48 The SQLAlchemy engine for the temporary database server. 

49 

50 Notes 

51 ----- 

52 This class was copied and modified from 

53 ``lsst.daf.butler.tests.postgresql``. 

54 """ 

55 

56 def __init__(self, server: Postgresql, engine: Engine) -> None: 

57 """Initialize the temporary Postgres database instance.""" 

58 self._server = server 

59 self._engine = engine 

60 

61 @property 

62 def url(self) -> str: 

63 """Return connection URL for the temporary database server. 

64 

65 Returns 

66 ------- 

67 str 

68 The connection URL. 

69 """ 

70 return self._server.url() 

71 

72 @property 

73 def engine(self) -> Engine: 

74 """Return the SQLAlchemy engine for the temporary database server. 

75 

76 Returns 

77 ------- 

78 `~sqlalchemy.engine.Engine` 

79 The SQLAlchemy engine. 

80 """ 

81 return self._engine 

82 

83 @contextmanager 

84 def begin(self) -> Iterator[Connection]: 

85 """Return a SQLAlchemy connection to the test database. 

86 

87 Returns 

88 ------- 

89 `~sqlalchemy.engine.Connection` 

90 The SQLAlchemy connection. 

91 """ 

92 with self._engine.begin() as connection: 

93 yield connection 

94 

95 def print_info(self) -> None: 

96 """Print information about the temporary database server.""" 

97 print("\n\n---- PostgreSQL URL ----") 

98 print(self.url) 

99 self._engine = create_engine(self.url) 

100 with self.begin() as conn: 

101 print("\n---- PostgreSQL Version ----") 

102 res = conn.execute(text("SELECT version()")).fetchone() 

103 if res: 

104 print(res[0]) 

105 print("\n") 

106 

107 

108@contextmanager 

109def setup_postgres_test_db() -> Iterator[TemporaryPostgresInstance]: 

110 """Set up a temporary Postgres database instance that can be used for 

111 testing. 

112 

113 Returns 

114 ------- 

115 TemporaryPostgresInstance 

116 The temporary Postgres database instance. 

117 

118 Raises 

119 ------ 

120 unittest.SkipTest 

121 Raised if the ``testing.postgresql`` module is not available. 

122 """ 

123 if Postgresql is None: 

124 raise unittest.SkipTest("testing.postgresql module not available.") 

125 

126 with Postgresql() as server: 

127 engine = create_engine(server.url()) 

128 instance = TemporaryPostgresInstance(server, engine) 

129 yield instance 

130 

131 # Clean up any lingering SQLAlchemy engines/connections 

132 # so they're closed before we shut down the server. 

133 engine.dispose() 

134 gc.collect()