Repository: apache/airflow
File: airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py (Python)

Code before:
            conn_id=TEST_CONN_ID,
            conn_type="sqlite",
            password="existing_password",
        )
        session.add(connection)
        session.commit()

        body = {
            "connection_id": TEST_CONN_ID,
            "conn_type": "sqlite",
            "password": "***",
        }
        response = test_client.post("/connections/test", json=body)
        assert response.status_code == 200
        assert response.json()["status"] is True

    @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
    def test_should_merge_extra_with_existing_connection(self, test_client, session):
        connection = Connection(
            conn_id=TEST_CONN_ID,
            conn_type="fs",
            extra='{"path": "/", "existing_key": "existing_value"}',
        )
        session.add(connection)
        session.commit()

        body = {
            "connection_id": TEST_CONN_ID,
            "conn_type": "fs",
            "extra": '{"path": "/", "new_key": "new_value"}',
        }
        response = test_client.post("/connections/test", json=body)
        assert response.status_code == 200
        assert response.json()["status"] is True

    @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
    def test_should_merge_both_password_and_extra(self, test_client, session):
        connection = Connection(
            conn_id=TEST_CONN_ID,
            conn_type="fs",
            password="existing_password",
            extra='{"path": "/", "existing_key": "existing_value"}',
        )
        session.add(connection)
        session.commit()

        body = {
            "connection_id": TEST_CONN_ID,
            "conn_type": "fs",
            "password": "***",
            "extra": '{"path": "/", "new_key": "new_value"}',

Review comment (suggestion):
I would also add an assertion to check that hte `TEST_CONN_ID` wasn't mutated into db, and that no other connection record was created. (I.e the merged connection is not persisted anywhere)

Code after:
            conn_id=TEST_CONN_ID,
            conn_type="sqlite",
            password="existing_password",
        )
        session.add(connection)
        session.commit()
        initial_count = session.scalar(select(func.count()).select_from(Connection))

        captured_value = {}

        def mock_test_connection(self):
            captured_value["password"] = self.password
            captured_value["conn_type"] = self.conn_type
            return True, "mocked"

        body = {
            "connection_id": TEST_CONN_ID,
            "conn_type": "new_sqlite",
            "password": "***",
        }

        with mock.patch.object(Connection, "test_connection", mock_test_connection):
            response = test_client.post("/connections/test", json=body)

        assert response.status_code == 200
        assert response.json()["status"] is True
        # Verify that the existing password was used, not "***"
        assert captured_value["password"] == "existing_password"
        # Verify that payload info were used for other fields
        assert captured_value["conn_type"] == "new_sqlite"

        # Verify DB was not mutated
        session.expire_all()
        db_conn = session.scalar(select(Connection).filter_by(conn_id=TEST_CONN_ID))
        assert db_conn.password == "existing_password"
        assert session.scalar(select(func.count()).select_from(Connection)) == initial_count

    @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"})
    def test_should_merge_extra_with_existing_connection(self, test_client, session):
        connection = Connection(
            conn_id=TEST_CONN_ID,
            conn_type="fs",
            extra='{"path": "/", "existing_key": "existing_value"}',
        )
        session.add(connection)
        session.commit()
        initial_count = session.scalar(select(func.count()).select_from(Connection))

        captured_extra = {}

        def mock_test_connection(self):