| 27 | |
| 28 | |
| 29 | def test_ssh_tunnel(mock_ssh_tunnel_forwarder: MagicMock, mock_pgexecute: MagicMock) -> None: |
| 30 | # Test with just a host |
| 31 | tunnel_url = "some.host" |
| 32 | db_params = { |
| 33 | "database": "dbname", |
| 34 | "host": "db.host", |
| 35 | "user": "db_user", |
| 36 | "passwd": "db_passwd", |
| 37 | } |
| 38 | expected_tunnel_params = { |
| 39 | "local_bind_address": ("127.0.0.1",), |
| 40 | "remote_bind_address": (db_params["host"], 5432), |
| 41 | "ssh_address_or_host": (tunnel_url, 22), |
| 42 | "logger": ANY, |
| 43 | } |
| 44 | |
| 45 | pgcli = PGCli(ssh_tunnel_url=tunnel_url) |
| 46 | pgcli.connect(**db_params) |
| 47 | |
| 48 | mock_ssh_tunnel_forwarder.assert_called_once_with(**expected_tunnel_params) |
| 49 | mock_ssh_tunnel_forwarder.return_value.start.assert_called_once() |
| 50 | mock_pgexecute.assert_called_once() |
| 51 | |
| 52 | call_args, call_kwargs = mock_pgexecute.call_args |
| 53 | assert call_args == ( |
| 54 | db_params["database"], |
| 55 | db_params["user"], |
| 56 | db_params["passwd"], |
| 57 | "127.0.0.1", |
| 58 | pgcli.ssh_tunnel.local_bind_ports[0], |
| 59 | "", |
| 60 | notify_callback, |
| 61 | ) |
| 62 | mock_ssh_tunnel_forwarder.reset_mock() |
| 63 | mock_pgexecute.reset_mock() |
| 64 | |
| 65 | # Test with a full url and with a specific db port |
| 66 | tunnel_user = "tunnel_user" |
| 67 | tunnel_passwd = "tunnel_pass" |
| 68 | tunnel_host = "some.other.host" |
| 69 | tunnel_port = 1022 |
| 70 | tunnel_url = f"ssh://{tunnel_user}:{tunnel_passwd}@{tunnel_host}:{tunnel_port}" |
| 71 | db_params["port"] = 1234 |
| 72 | |
| 73 | expected_tunnel_params["remote_bind_address"] = ( |
| 74 | db_params["host"], |
| 75 | db_params["port"], |
| 76 | ) |
| 77 | expected_tunnel_params["ssh_address_or_host"] = (tunnel_host, tunnel_port) |
| 78 | expected_tunnel_params["ssh_username"] = tunnel_user |
| 79 | expected_tunnel_params["ssh_password"] = tunnel_passwd |
| 80 | |
| 81 | pgcli = PGCli(ssh_tunnel_url=tunnel_url) |
| 82 | pgcli.connect(**db_params) |
| 83 | |
| 84 | mock_ssh_tunnel_forwarder.assert_called_once_with(**expected_tunnel_params) |
| 85 | mock_ssh_tunnel_forwarder.return_value.start.assert_called_once() |
| 86 | mock_pgexecute.assert_called_once() |