Coverage for chatgpt_proxy / tests / setup.py: 83%

1# MIT License 

2# 

3# Copyright (c) 2025 Tuomo Kriikkula 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23import datetime 

24import os 

25import time 

26from collections.abc import Callable 

27from contextlib import contextmanager 

28from pathlib import Path 

29from typing import Any 

30from urllib.parse import urlparse 

31 

32import asyncpg 

33 

34import chatgpt_proxy 

35 

36 

37class Ignored: 

38 pass 

39 

40 

41IGNORED = Ignored() 

42 

43test_sanic_secret = "dummy123123-45678910111213141516" 

44test_db = "chatgpt_proxy_tests" 

45default_db_url = "postgresql://postgres:postgres@localhost:5432" 

46db_url = default_db_url 

47db_base_url = default_db_url 

48db_test_url = f"{db_url.rstrip("/")}/chatgpt_proxy_tests" 

49steam_web_api_key = "dummy_steam_web_api_key" 

50 

51_root_path = Path(chatgpt_proxy.__file__).resolve().parent 

52_pkg_path_db = _root_path / "db/" 

53_pkg_path_tests = _root_path / "tests/" 

54 

55assert _root_path.exists() 

56assert _pkg_path_db.exists() 

57assert _pkg_path_tests.exists() 

58 

59default_test_db_timeout = 30.0 

60 

61 

62async def initialize_test_db( 

63 conn: asyncpg.Connection, 

64 timeout: float | None = default_test_db_timeout 

65): 

66 init_db_sql = (_pkg_path_db / "db.sql").read_text() 

67 await conn.execute(init_db_sql, timeout=timeout) 

68 

69 

70async def seed_test_db( 

71 conn: asyncpg.Connection, 

72 timeout: float | None = default_test_db_timeout 

73): 

74 seed_db_sql = (_pkg_path_tests / "seed.sql").read_text() 

75 await conn.execute(seed_db_sql, timeout=timeout) 

76 

77 

78async def drop_test_db( 

79 conn: asyncpg.Connection, 

80 timeout: float | None = default_test_db_timeout 

81): 

82 await conn.execute( 

83 f"DROP DATABASE IF EXISTS {test_db} WITH (FORCE)", 

84 timeout=timeout, 

85 ) 

86 

87 

88async def create_test_db( 

89 conn: asyncpg.Connection, 

90 timeout: float | None = default_test_db_timeout 

91): 

92 await conn.execute( 

93 f"CREATE DATABASE {test_db};", 

94 timeout=timeout, 

95 ) 

96 

97 

98# TODO: this is sketchy as fuck, try to come up with a better way? 

99def common_test_setup( 

100 db_url_: str | Ignored = IGNORED, 

101 test_db_: str | Ignored = IGNORED, 

102 db_test_url_: str | Ignored = IGNORED, 

103 steam_web_api_key_: str | Ignored = IGNORED, 

104 test_sanic_secret_: str | Ignored = IGNORED, 

105) -> None: 

106 global db_url 

107 global db_base_url 

108 global db_test_url 

109 global test_db 

110 global steam_web_api_key 

111 global test_sanic_secret 

112 

113 if db_url_ is not IGNORED: 

114 db_url = str(db_url_) 

115 else: 

116 # NOTE: avoid messing up DB URL if this is called multiple times 

117 # from different tests! 

118 # TODO: make a cleaner way of handling this! 

119 db_url = os.environ.get("DATABASE_URL", default_db_url) 

120 parts = urlparse(db_url) 

121 parts = parts._replace(path="") 

122 db_base_url = parts.geturl() 

123 

124 if db_test_url_ is not IGNORED: 

125 db_test_url = str(db_test_url_) 

126 else: 

127 # NOTE: avoid messing up DB URL if this is called multiple times 

128 # from different tests! 

129 # TODO: make a cleaner way of handling this! 

130 db_test_url = f"{db_base_url.rstrip("/")}/chatgpt_proxy_tests" 

131 

132 if test_db_ is not IGNORED: 

133 test_db = str(test_db_) 

134 if steam_web_api_key_ is not IGNORED: 

135 steam_web_api_key = str(steam_web_api_key_) 

136 if test_sanic_secret_ is not IGNORED: 

137 test_sanic_secret = str(test_sanic_secret_) 

138 

139 os.environ["SANIC_SECRET"] = test_sanic_secret 

140 os.environ["OPENAI_API_KEY"] = "dummy" 

141 os.environ["DATABASE_URL"] = db_test_url 

142 os.environ["STEAM_WEB_API_KEY"] = steam_web_api_key 

143 

144 

145@contextmanager 

146def retry_context( 

147 builder: Callable[[], Any], 

148 retries: int = 5, 

149 delay: datetime.timedelta | None = None, 

150 exc_types: tuple[type[Exception]] | None = None, 

151 retry_cb: Callable[[Exception, int], None] | None = None, 

152) -> Any: 

153 if delay is None: 

154 delay = datetime.timedelta(milliseconds=100) 

155 if exc_types is None: 

156 exc_types = (Exception,) 

157 

158 retry = 1 

159 while True: 

160 try: 

161 with builder() as obj: 

162 retries += 1 

163 yield obj 

164 return 

165 except exc_types as e: 

166 if retry_cb: 

167 retry_cb(e, retry) 

168 if retry > retries: 

169 raise 

170 time.sleep(delay.total_seconds())