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