Coverage for chatgpt_proxy / tests / test_gen_api_key.py: 100%

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 hashlib 

25import ipaddress 

26from typing import AsyncGenerator 

27 

28import asyncpg 

29import pytest 

30import pytest_asyncio 

31 

32from chatgpt_proxy.db import pool_acquire 

33from chatgpt_proxy.db import queries 

34from chatgpt_proxy.gen_api_key import async_main 

35from chatgpt_proxy.tests import setup 

36from chatgpt_proxy.tests.setup import default_test_db_timeout 

37 

38setup.common_test_setup() 

39 

40_db_timeout = default_test_db_timeout 

41 

42 

43@pytest_asyncio.fixture 

44async def gen_api_key_fixture() -> AsyncGenerator[asyncpg.Connection]: 

45 db_fixture_pool = await asyncpg.create_pool( 

46 dsn=setup.db_base_url, 

47 min_size=1, 

48 max_size=1, 

49 timeout=_db_timeout, 

50 ) 

51 

52 async with pool_acquire(db_fixture_pool, timeout=_db_timeout) as conn: 

53 await setup.drop_test_db(conn, timeout=_db_timeout) 

54 await setup.create_test_db(conn, timeout=_db_timeout) 

55 

56 test_db_pool = await asyncpg.create_pool( 

57 dsn=setup.db_test_url, 

58 min_size=1, 

59 max_size=1, 

60 timeout=_db_timeout, 

61 ) 

62 

63 async with pool_acquire( 

64 test_db_pool, 

65 timeout=_db_timeout, 

66 ) as conn: 

67 async with conn.transaction(): 

68 await setup.initialize_test_db(conn, timeout=_db_timeout) 

69 

70 yield conn 

71 

72 async with pool_acquire(db_fixture_pool, timeout=_db_timeout) as conn: 

73 await setup.drop_test_db(conn, timeout=_db_timeout) 

74 

75 await db_fixture_pool.close() 

76 await test_db_pool.close() 

77 

78 

79@pytest.mark.asyncio 

80async def test_gen_api_key_async_main(gen_api_key_fixture): 

81 conn = gen_api_key_fixture 

82 

83 token = await async_main( 

84 game_server_address=ipaddress.IPv4Address("69.69.69.1"), 

85 game_server_port=60000, 

86 secret=setup.test_sanic_secret, 

87 issuer="test123", 

88 audience="test123", 

89 expires_at=datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(days=1), 

90 name="test_key_blabla", 

91 ) 

92 

93 assert token 

94 

95 key = await queries.select_game_server_api_key( 

96 conn=conn, 

97 game_server_address=ipaddress.IPv4Address("69.69.69.1"), 

98 game_server_port=60000, 

99 ) 

100 assert key 

101 

102 assert key["api_key_hash"] == hashlib.sha256(token.encode()).digest()