|
18 | 18 | import contextlib |
19 | 19 | from datetime import datetime |
20 | 20 | from datetime import timezone |
| 21 | +import logging |
21 | 22 | import os |
22 | 23 | import pickle |
23 | 24 | import time |
| 25 | +from unittest import mock |
24 | 26 |
|
25 | 27 | from fastapi.openapi.models import HTTPBearer |
26 | 28 | from google.adk.auth.auth_tool import AuthConfig |
|
29 | 31 | from google.adk.events.ui_widget import UiWidget |
30 | 32 | from google.adk.sessions.migration import _schema_check_utils |
31 | 33 | from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle as mfsp |
| 34 | +from google.adk.sessions.migration import migrate_from_sqlalchemy_sqlite as mfss |
| 35 | +from google.adk.sessions.migration import migration_runner |
32 | 36 | from google.adk.sessions.schemas import v0 |
33 | 37 | from google.adk.sessions.schemas import v1 |
34 | 38 | from google.adk.tools.tool_confirmation import ToolConfirmation |
@@ -116,6 +120,118 @@ def test_to_sync_url_empty_string(self): |
116 | 120 | assert _schema_check_utils.to_sync_url("") == "" |
117 | 121 |
|
118 | 122 |
|
| 123 | +class TestRedactDbUrl: |
| 124 | + """Tests for the _redact_db_url function.""" |
| 125 | + |
| 126 | + def test_password_is_masked(self): |
| 127 | + redacted = _schema_check_utils._redact_db_url( |
| 128 | + "postgresql+asyncpg://user:sup3r-s3cret@host:5432/db" |
| 129 | + ) |
| 130 | + assert redacted == "postgresql+asyncpg://user:***@host:5432/db" |
| 131 | + |
| 132 | + def test_unparseable_url_falls_back_to_placeholder(self): |
| 133 | + """Redaction runs while reporting an error, so it must never raise.""" |
| 134 | + assert ( |
| 135 | + _schema_check_utils._redact_db_url("definitely not a url sup3r-s3cret") |
| 136 | + == "<unparseable database URL>" |
| 137 | + ) |
| 138 | + |
| 139 | + def test_query_parameter_values_are_masked(self): |
| 140 | + """Drivers accept secrets as query parameters, so every value is masked.""" |
| 141 | + redacted = _schema_check_utils._redact_db_url( |
| 142 | + "postgresql://user@host:5432/db?password=sup3r-s3cret&sslmode=require" |
| 143 | + ) |
| 144 | + assert redacted == ( |
| 145 | + "postgresql://user@host:5432/db?password=REDACTED&sslmode=REDACTED" |
| 146 | + ) |
| 147 | + |
| 148 | + def test_schema_version_failure_warning_hides_password(self, caplog): |
| 149 | + db_url = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/db" |
| 150 | + |
| 151 | + with mock.patch.object( |
| 152 | + _schema_check_utils, |
| 153 | + "create_sync_engine", |
| 154 | + side_effect=RuntimeError("boom"), |
| 155 | + ): |
| 156 | + with caplog.at_level(logging.WARNING): |
| 157 | + with pytest.raises(RuntimeError): |
| 158 | + _schema_check_utils.get_db_schema_version(db_url) |
| 159 | + |
| 160 | + assert "sup3r-s3cret" not in caplog.text |
| 161 | + assert "postgresql+asyncpg://user:***@host:5432/db" in caplog.text |
| 162 | + |
| 163 | + |
| 164 | +_SOURCE_URL = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/src" |
| 165 | +_DEST_URL = "postgresql+asyncpg://user:0ther-s3cret@host:5432/dst" |
| 166 | + |
| 167 | + |
| 168 | +class TestMigrationLogsHidePassword: |
| 169 | + """These entry points log their URLs on every run, not only on failure.""" |
| 170 | + |
| 171 | + def test_pickle_migration_connect_logs_are_redacted(self, caplog): |
| 172 | + with mock.patch.object( |
| 173 | + mfsp, |
| 174 | + "create_engine", |
| 175 | + side_effect=[mock.MagicMock(), RuntimeError("boom")], |
| 176 | + ): |
| 177 | + with caplog.at_level(logging.INFO): |
| 178 | + with pytest.raises(RuntimeError): |
| 179 | + mfsp.migrate(_SOURCE_URL, _DEST_URL) |
| 180 | + |
| 181 | + assert "sup3r-s3cret" not in caplog.text |
| 182 | + assert "0ther-s3cret" not in caplog.text |
| 183 | + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text |
| 184 | + assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text |
| 185 | + |
| 186 | + def test_sqlite_migration_connect_log_is_redacted(self, caplog, tmp_path): |
| 187 | + with mock.patch.object( |
| 188 | + mfss, "create_engine", side_effect=RuntimeError("boom") |
| 189 | + ): |
| 190 | + with caplog.at_level(logging.INFO): |
| 191 | + with pytest.raises(SystemExit): |
| 192 | + mfss.migrate(_SOURCE_URL, str(tmp_path / "dest.db")) |
| 193 | + |
| 194 | + assert "sup3r-s3cret" not in caplog.text |
| 195 | + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text |
| 196 | + |
| 197 | + def test_runner_up_to_date_log_is_redacted(self, caplog): |
| 198 | + with mock.patch.object( |
| 199 | + _schema_check_utils, |
| 200 | + "get_db_schema_version", |
| 201 | + return_value=migration_runner.LATEST_VERSION, |
| 202 | + ): |
| 203 | + with caplog.at_level(logging.INFO): |
| 204 | + migration_runner.upgrade(_SOURCE_URL, _DEST_URL) |
| 205 | + |
| 206 | + assert "sup3r-s3cret" not in caplog.text |
| 207 | + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text |
| 208 | + |
| 209 | + def test_runner_migration_step_log_is_redacted(self, caplog): |
| 210 | + mock_migrate = mock.Mock() |
| 211 | + with mock.patch.object( |
| 212 | + _schema_check_utils, |
| 213 | + "get_db_schema_version", |
| 214 | + return_value=_schema_check_utils.SCHEMA_VERSION_0_PICKLE, |
| 215 | + ): |
| 216 | + with mock.patch.dict( |
| 217 | + migration_runner.MIGRATIONS, |
| 218 | + { |
| 219 | + _schema_check_utils.SCHEMA_VERSION_0_PICKLE: ( |
| 220 | + _schema_check_utils.SCHEMA_VERSION_1_JSON, |
| 221 | + mock_migrate, |
| 222 | + ) |
| 223 | + }, |
| 224 | + ): |
| 225 | + with caplog.at_level(logging.INFO): |
| 226 | + migration_runner.upgrade(_SOURCE_URL, _DEST_URL) |
| 227 | + |
| 228 | + mock_migrate.assert_called_once_with(_SOURCE_URL, _DEST_URL) |
| 229 | + assert "sup3r-s3cret" not in caplog.text |
| 230 | + assert "0ther-s3cret" not in caplog.text |
| 231 | + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text |
| 232 | + assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text |
| 233 | + |
| 234 | + |
119 | 235 | def test_migrate_from_sqlalchemy_pickle(tmp_path): |
120 | 236 | """Tests for migrate_from_sqlalchemy_pickle.""" |
121 | 237 | source_db_path = tmp_path / "source_pickle.db" |
|
0 commit comments