feat: default ports for SSH tunnel (#32403)

tests/unit_tests/databases/ssh_tunnel/commands/create_test.py

@@ -17,6 +17,7 @@
 
 
 import pytest
+from sqlalchemy.orm.session import Session
 
 from superset.commands.database.ssh_tunnel.exceptions import (
     SSHTunnelDatabasePortError,
@@ -24,13 +25,16 @@ from superset.commands.database.ssh_tunnel.exceptions import (
 )
 
 
-def test_create_ssh_tunnel_command() -> None:
+def test_create_ssh_tunnel_command(session: Session) -> None:
+    from superset import db
     from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
     from superset.databases.ssh_tunnel.models import SSHTunnel
     from superset.models.core import Database
 
+    engine = db.session.get_bind()
+    Database.metadata.create_all(engine)  # pylint: disable=no-member
+
     database = Database(
-        id=1,
         database_name="my_database",
         sqlalchemy_uri="postgresql://u:p@localhost:5432/db",
     )
@@ -49,12 +53,15 @@ def test_create_ssh_tunnel_command() -> None:
     assert isinstance(result, SSHTunnel)
 
 
-def test_create_ssh_tunnel_command_invalid_params() -> None:
+def test_create_ssh_tunnel_command_invalid_params(session: Session) -> None:
+    from superset import db
     from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
     from superset.models.core import Database
 
+    engine = db.session.get_bind()
+    Database.metadata.create_all(engine)  # pylint: disable=no-member
+
     database = Database(
-        id=1,
         database_name="my_database",
         sqlalchemy_uri="postgresql://u:p@localhost:5432/db",
     )
@@ -76,14 +83,52 @@ def test_create_ssh_tunnel_command_invalid_params() -> None:
     assert str(excinfo.value) == ("SSH Tunnel parameters are invalid.")
 
 
-def test_create_ssh_tunnel_command_no_port() -> None:
+def test_create_ssh_tunnel_command_no_port(session: Session) -> None:
+    """
+    Test that SSH Tunnel can be created without explicit port but with a default one.
+    """
+    from superset import db
+    from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
+    from superset.databases.ssh_tunnel.models import SSHTunnel
+    from superset.models.core import Database
+
+    engine = db.session.get_bind()
+    Database.metadata.create_all(engine)  # pylint: disable=no-member
+
+    database = Database(
+        database_name="my_database",
+        sqlalchemy_uri="postgresql://u:p@localhost/db",
+    )
+
+    properties = {
+        "database": database,
+        "server_address": "123.132.123.1",
+        "server_port": "3005",
+        "username": "foo",
+        "password": "bar",
+    }
+
+    result = CreateSSHTunnelCommand(database, properties).run()
+
+    assert result is not None
+    assert isinstance(result, SSHTunnel)
+
+
+def test_create_ssh_tunnel_command_no_port_no_default(session: Session) -> None:
+    """
+    Test that error is raised when creating SSH Tunnel without explicit/default ports.
+    """
+    from superset import db
     from superset.commands.database.ssh_tunnel.create import CreateSSHTunnelCommand
     from superset.models.core import Database
 
+    engine = db.session.get_bind()
+    Database.metadata.create_all(engine)  # pylint: disable=no-member
+
     database = Database(
         id=1,
         database_name="my_database",
-        sqlalchemy_uri="postgresql://u:p@localhost/db",
+        sqlalchemy_uri="weird+db://u:p@localhost/db",
     )
 
     properties = {

tests/unit_tests/databases/ssh_tunnel/commands/update_test.py

@@ -103,6 +103,37 @@ def test_update_shh_tunnel_invalid_params(session_with_data: Session) -> None:
     "session_with_data", ["postgresql://u:p@localhost/testdb"], indirect=True
 )
 def test_update_shh_tunnel_no_port(session_with_data: Session) -> None:
+    """
+    Test that SSH Tunnel can be updated without explicit port but with a default one.
+    """
+    from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand
+    from superset.daos.database import DatabaseDAO
+    from superset.databases.ssh_tunnel.models import SSHTunnel
+
+    result = DatabaseDAO.get_ssh_tunnel(1)
+
+    assert result
+    assert isinstance(result, SSHTunnel)
+    assert 1 == result.database_id
+    assert "Test" == result.server_address
+
+    update_payload = {"server_address": "Test2"}
+    UpdateSSHTunnelCommand(1, update_payload).run()
+
+    result = DatabaseDAO.get_ssh_tunnel(1)
+
+    assert result
+    assert isinstance(result, SSHTunnel)
+    assert "Test2" == result.server_address
+
+
+@pytest.mark.parametrize(
+    "session_with_data", ["weird+db://u:p@localhost/testdb"], indirect=True
+)
+def test_update_shh_tunnel_no_port_no_default(session_with_data: Session) -> None:
+    """
+    Test that error is raised when updating SSH Tunnel without explicit/default ports.
+    """
     from superset.commands.database.ssh_tunnel.update import UpdateSSHTunnelCommand
     from superset.daos.database import DatabaseDAO
     from superset.databases.ssh_tunnel.models import SSHTunnel