Python OSLoginHook 예제들

예제 #1

파일: test_os_login.py 프로젝트: 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,
        )

예제 #2

파일: test_os_login.py 프로젝트: 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,
            )

예제 #3

파일: test_os_login.py 프로젝트: phoenix-elite1050/airflow

 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")

예제 #4

파일: compute_ssh.py 프로젝트: thesuperzapper/airflow

 def _oslogin_hook(self) -> OSLoginHook:
     return OSLoginHook(gcp_conn_id=self.gcp_conn_id,
                        delegate_to=self.delegate_to)