Coverage for app/tests/alembic/test_utils.py: 98%

1from pathlib import Path 

2from unittest import mock 

3 

4import pytest 

5from sqlalchemy import Engine, text 

6 

7from app.alembic.utils import ( 

8 get_sql_migration_file, 

9 read_sql_file, 

10 run_sql_statements, 

11) 

12 

13 

14@pytest.fixture(scope="session") 

15def project_root() -> Path: 

16 return Path(__file__).parents[3] 

17 

18 

19@pytest.fixture 

20def alembic_migration_file(project_root: Path) -> Path: 

21 return project_root / "app/alembic/versions/eb3ad385519a_init_database.py" 

22 

23 

24@pytest.fixture 

25def alembic_sql_upgrade_file(project_root: Path) -> Path: 

26 return project_root / "app/alembic/versions/eb3ad385519a_init_database_upgrade.sql" 

27 

28 

29@pytest.fixture 

30def alembic_sql_downgrade_file(project_root: Path) -> Path: 

31 return ( 

32 project_root / "app/alembic/versions/eb3ad385519a_init_database_downgrade.sql" 

33 ) 

34 

35 

36def test_read_sql_file_missing() -> None: 

37 with pytest.raises(FileNotFoundError): 

38 read_sql_file(Path("/tmp/file.sql")) 

39 

40 

41def test_read_sql_file(alembic_sql_upgrade_file: Path) -> None: 

42 def normalize_whitespace(sql: str) -> str: 

43 return " ".join(sql.split()) 

44 

45 result = tuple( 

46 normalize_whitespace(stmt) for stmt in read_sql_file(alembic_sql_upgrade_file) 

47 ) 

48 expected_result = tuple( 

49 normalize_whitespace(stmt) 

50 for stmt in ( 

51 "pragma journal_mode = wal", 

52 """create table pypi_package_downloads_weekly_metrics ( 

53 package_name text not null, 

54 package_downloaded_week text not null, 

55 downloads integer not null, 

56 cumulative_downloads integer not null, 

57 first_distribution_week text not null, 

58 weeks_since_first_distribution integer not null, 

59 synced_at text not null, 

60 primary key (package_name, package_downloaded_week) 

61 ) 

62 """, 

63 """create table pypi_packages ( 

64 package_name text not null primary key, 

65 latest_package_version text not null, 

66 package_summary text, 

67 package_home_page text, 

68 package_download_url text, 

69 package_uploaded_at text not null, 

70 synced_at text not null 

71 ) 

72 """, 

73 ) 

74 ) 

75 assert result == expected_result 

76 

77 

78def test_get_sql_file_name_upgrade( 

79 alembic_migration_file: Path, monkeypatch: pytest.MonkeyPatch 

80) -> None: 

81 mock_stack = [ 

82 mock.Mock(), 

83 mock.Mock(filename=str(alembic_migration_file), function="upgrade"), 

84 ] 

85 monkeypatch.setattr("inspect.stack", lambda: mock_stack) 

86 result = get_sql_migration_file() 

87 excepted_result = ( 

88 alembic_migration_file.parent / f"{alembic_migration_file.stem}_upgrade.sql" 

89 ) 

90 assert result == excepted_result 

91 

92 

93def test_get_sql_file_name_downgrade( 

94 alembic_migration_file: Path, monkeypatch: pytest.MonkeyPatch 

95) -> None: 

96 mock_stack = [ 

97 mock.Mock(), 

98 mock.Mock(filename=str(alembic_migration_file), function="downgrade"), 

99 ] 

100 monkeypatch.setattr("inspect.stack", lambda: mock_stack) 

101 result = get_sql_migration_file() 

102 excepted_result = ( 

103 alembic_migration_file.parent / f"{alembic_migration_file.stem}_downgrade.sql" 

104 ) 

105 assert result == excepted_result 

106 

107 

108def test_get_sql_file_name_invalid( 

109 alembic_migration_file: Path, monkeypatch: pytest.MonkeyPatch 

110) -> None: 

111 mock_stack = [ 

112 mock.Mock(), 

113 mock.Mock(filename=str(alembic_migration_file), function="invalid"), 

114 ] 

115 monkeypatch.setattr("inspect.stack", lambda: mock_stack) 

116 with pytest.raises( 

117 ValueError, 

118 match="Function 'invalid' is not a valid Alembic migration function. Expected 'upgrade' or 'downgrade'.", 

119 ): 

120 get_sql_migration_file() 

121 

122 

123def test_run_sql_statements( 

124 write_engine: Engine, read_engine: Engine, tmp_path: Path 

125) -> None: 

126 sql_file = tmp_path / "test_statements.sql" 

127 sql_file.write_text( 

128 """ 

129 CREATE TABLE test_table ( 

130 id INTEGER PRIMARY KEY, 

131 name TEXT NOT NULL 

132 ); 

133 INSERT INTO test_table (id, name) VALUES (1, 'test_name'); 

134 """ 

135 ) 

136 

137 executed_statements = run_sql_statements(write_engine, sql_file) 

138 

139 assert len(executed_statements) == 2 

140 assert executed_statements[0].startswith("CREATE TABLE test_table") 

141 assert executed_statements[1].startswith("INSERT INTO test_table") 

142 

143 with read_engine.connect() as conn: 

144 result = conn.execute(text("SELECT id, name FROM test_table")).fetchall() 

145 

146 assert result == [(1, "test_name")] 

147 

148 

149def test_run_sql_statements_error(write_engine: Engine, tmp_path: Path) -> None: 

150 sql_file = tmp_path / "test_invalid_statements.sql" 

151 sql_file.write_text("INSERT INTO non_existent_table (id) VALUES (1);") 

152 

153 with pytest.raises(Exception): 

154 run_sql_statements(write_engine, sql_file)