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

1"""Module for testing the data transfer objects TableRow and TableJoinRow.""" 

2 

3from dataclasses import dataclass 

4 

5import pytest 

6 

7from kwai_core.db.table_row import JoinedTableRow, TableRow 

8from sql_smith.engine import CommonEngine 

9 

10 

11@dataclass(kw_only=True, frozen=True, slots=True) 

12class JudokaRow(TableRow): 

13 """Data transfer object for a row of the judokas table.""" 

14 

15 __table_name__ = "judokas" 

16 

17 id: int 

18 name: str 

19 age: int 

20 

21 

22def test_aliases(): 

23 """Test if the aliases are created correctly.""" 

24 aliases = JudokaRow.get_aliases() 

25 assert len(aliases) == 3, "There should be 3 aliases." 

26 assert aliases[0].sql(CommonEngine()) == '"judokas"."id" AS "judokas_id"' 

27 assert aliases[1].sql(CommonEngine()) == '"judokas"."name" AS "judokas_name"' 

28 assert aliases[2].sql(CommonEngine()) == '"judokas"."age" AS "judokas_age"' 

29 

30 

31def test_column(): 

32 """Test if the column is created correctly.""" 

33 column = JudokaRow.column("name") 

34 assert column == "judokas.name" 

35 

36 

37def test_create_aliases_with_other_table_name(): 

38 """Test the creation of column aliases when the table name is also aliased.""" 

39 aliases = JudokaRow.get_aliases("my_judokas") 

40 assert len(aliases) == 3, "There should be 3 aliases" 

41 assert aliases[0].sql(CommonEngine()) == '"judokas"."id" AS "my_judokas_id"' 

42 assert aliases[1].sql(CommonEngine()) == '"judokas"."name" AS "my_judokas_name"' 

43 assert aliases[2].sql(CommonEngine()) == '"judokas"."age" AS "my_judokas_age"' 

44 

45 

46def test_map_row(): 

47 """Test map row.""" 

48 row = JudokaRow.map({"judokas_id": 1, "judokas_name": "Jigoro", "judokas_age": 77}) 

49 assert row == JudokaRow(id=1, name="Jigoro", age=77) 

50 

51 

52def test_invalid_map_row(): 

53 """Test map row.""" 

54 with pytest.raises(ValueError, match=r"^id\(1\) of JudokaRow"): 

55 JudokaRow.map({"judokas_id": "1", "judokas_name": "Jigoro", "judokas_age": 77}) 

56 

57 

58@dataclass(kw_only=True, frozen=True, slots=True) 

59class CountryRow(TableRow): 

60 """Data transfer object for a row of the countries table.""" 

61 

62 __table_name__ = "countries" 

63 

64 iso_2: str 

65 name: str 

66 

67 

68@dataclass(kw_only=True, frozen=True, slots=True) 

69class MemberRow(JoinedTableRow): 

70 """A data transfer object for a query that joins multiple tables..""" 

71 

72 judoka: JudokaRow 

73 country: CountryRow 

74 

75 

76def test_joined_table_row(): 

77 """Test JoinedTableRow.""" 

78 aliases = MemberRow.get_aliases() 

79 assert len(aliases) == 5, "There should be 5 aliases." 

80 assert aliases[0].sql(CommonEngine()) == '"judokas"."id" AS "judoka_id"' 

81 assert aliases[1].sql(CommonEngine()) == '"judokas"."name" AS "judoka_name"' 

82 assert aliases[2].sql(CommonEngine()) == '"judokas"."age" AS "judoka_age"' 

83 assert aliases[3].sql(CommonEngine()) == '"countries"."iso_2" AS "country_iso_2"' 

84 assert aliases[4].sql(CommonEngine()) == '"countries"."name" AS "country_name"' 

85 

86 

87def test_map_joined_tables(): 

88 """Test mapping from a row with joined tables.""" 

89 row = MemberRow.map( 

90 { 

91 "judoka_id": 1, 

92 "judoka_name": "Jigoro", 

93 "judoka_age": 77, 

94 "country_iso_2": "JP", 

95 "country_name": "Japan", 

96 } 

97 ) 

98 assert row.judoka == JudokaRow(id=1, name="Jigoro", age=77) 

99 assert row.country == CountryRow(iso_2="JP", name="Japan")