Coverage for src/tests/core/db/test_database.py: 100%

1"""Module for testing Database class.""" 

2 

3import pytest 

4 

5from kwai_core.db.database import Database 

6from sql_smith.functions import alias, func 

7from sql_smith.query import SelectQuery 

8 

9 

10pytestmark = pytest.mark.db 

11 

12 

13async def test_fetch(database: Database): 

14 """Test fetch.""" 

15 select: SelectQuery = ( 

16 Database.create_query_factory() 

17 .select(alias(func("COUNT", func("DISTINCT", "email")), "c")) 

18 .from_("users") 

19 ) 

20 record = await database.fetch_one(select) 

21 assert record is not None, "There should be at least a record" 

22 assert "c" in record, "There should be a key 'c' in the record"