Exemplos de OSLoginHook.import_ssh_public_key em Python

Exemplo n.º 1

Arquivo: test_os_login.py Projeto: phoenix-elite1050/airflow

class TestOSLoginHookWithDefaultProjectIdHook(TestCase):
    def setUp(
        self,
    ) -> None:
        with mock.patch(
            "airflow.providers.google.cloud.hooks.os_login.OSLoginHook.__init__",
            new=mock_base_gcp_hook_default_project_id,
        ):
            self.hook = OSLoginHook(gcp_conn_id="test")

    @mock.patch(
        "airflow.providers.google.common.hooks.base_google.GoogleBaseHook._get_credentials_and_project_id",
        return_value=(TEST_CREDENTIALS, TEST_PROJECT_ID_2),
    )
    @mock.patch("airflow.providers.google.cloud.hooks.os_login.OSLoginHook.get_conn")
    def test_import_ssh_public_key(self, mock_get_conn, mock_get_creds_and_project_id) -> None:
        self.hook.import_ssh_public_key(
            user=TEST_USER,
            ssh_public_key=TEST_BODY,
            project_id=None,
            retry=TEST_RETRY,
            timeout=TEST_TIMEOUT,
            metadata=TEST_METADATA,
        )
        mock_get_conn.return_value.import_ssh_public_key.assert_called_once_with(
            parent=TEST_PARENT,
            ssh_public_key=TEST_BODY,
            project_id=TEST_PROJECT_ID_2,
            retry=TEST_RETRY,
            timeout=TEST_TIMEOUT,
            metadata=TEST_METADATA,
        )

Exemplo n.º 2

Arquivo: test_os_login.py Projeto: phoenix-elite1050/airflow

class TestOSLoginHookMissingProjectIdHook(TestCase):
    def setUp(
        self,
    ) -> None:
        with mock.patch(
            "airflow.providers.google.cloud.hooks.os_login.OSLoginHook.__init__",
            new=mock_base_gcp_hook_no_default_project_id,
        ):
            self.hook = OSLoginHook(gcp_conn_id="test")

    @mock.patch(
        "airflow.providers.google.common.hooks.base_google.GoogleBaseHook._get_credentials_and_project_id",
        return_value=(TEST_CREDENTIALS, None),
    )
    @mock.patch("airflow.providers.google.cloud.hooks.os_login.OSLoginHook.get_conn")
    def test_import_ssh_public_key(self, mock_get_conn, mock_get_creds_and_project_id) -> None:
        with self.assertRaisesRegex(AirflowException, TEST_MESSAGE):
            self.hook.import_ssh_public_key(
                user=TEST_USER,
                ssh_public_key=TEST_BODY,
                project_id=None,
                retry=TEST_RETRY,
                timeout=TEST_TIMEOUT,
                metadata=TEST_METADATA,
            )