Mock all connections in `TestYandexCloudYqHook::test_select_results`

[Taragolis]

trafficstars

Body

Originally reported in Slack

Jed Cunningham 7:23 AM

Feels like we might have missed mocking something in the yandex test suite. We are getting grpc failures - seems their service might be having troubles? Either way, I can’t imagine we really want our test suite doing grpc anyways. https://github.com/apache/airflow/actions/runs/8904604506/job/24454208443

__________________ TestYandexCloudYqHook.test_select_results ___________________
[gw3] linux -- Python 3.11.9 /usr/local/bin/python

self = <tests.providers.yandex.hooks.test_yq.TestYandexCloudYqHook object at 0x7f65efadec90>
mock_get_auth_token_requester = <MagicMock name='get_auth_token_requester' id='140075695241040'>
mock_validate = <MagicMock name='__validate_service_account_key' id='140075695229520'>
mock_create_token = <PropertyMock name='Create' id='140075529000656'>

    @mock.patch(
        "yandex.cloud.iam.v1.iam_token_service_pb2_grpc.IamTokenServiceStub.Create",
        create=True,
        new_callable=mock.PropertyMock,
    )
    @mock.patch("yandexcloud._auth_fabric.__validate_service_account_key")
    @mock.patch("yandexcloud._auth_fabric.get_auth_token_requester", return_value=DummyTokenRequester())
    def test_select_results(self, mock_get_auth_token_requester, mock_validate, mock_create_token):
        with mock.patch.multiple(
            "yandex_query_client.YQHttpClient",
            create_query=mock.DEFAULT,
            wait_query_to_succeed=mock.DEFAULT,
            get_query_all_result_sets=mock.DEFAULT,
            get_query_status=mock.DEFAULT,
            get_query=mock.DEFAULT,
            stop_query=mock.DEFAULT,
        ) as mocks:
>           self._init_hook()

tests/providers/yandex/hooks/test_yq.py:121: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/providers/yandex/hooks/test_yq.py:55: in _init_hook
    self.hook = YQHook(default_folder_id="my_folder_id")
airflow/providers/yandex/hooks/yq.py:38: in __init__
    token=self._get_iam_token(), project=self.default_folder_id, user_agent=provider_user_agent()
airflow/providers/yandex/hooks/yq.py:111: in _get_iam_token
    client = sdk.client(IamTokenServiceStub)
/usr/local/lib/python3.11/site-packages/yandexcloud/_sdk.py:42: in client
    channel = self._channels.channel(service)
/usr/local/lib/python3.11/site-packages/yandexcloud/_channels.py:50: in channel
    resp = endpoint_service.List(ListApiEndpointsRequest())
/usr/local/lib/python3.11/site-packages/grpc/_channel.py:1181: in __call__
    return _end_unary_response_blocking(state, call, False, None)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

state = <grpc._channel._RPCState object at 0x7f65e9e09950>
call = <grpc._cython.cygrpc.SegregatedCall object at 0x7f65e025ce80>
with_call = False, deadline = None

    def _end_unary_response_blocking(
        state: _RPCState,
        call: cygrpc.SegregatedCall,
        with_call: bool,
        deadline: Optional[float],
    ) -> Union[ResponseType, Tuple[ResponseType, grpc.Call]]:
        if state.code is grpc.StatusCode.OK:
            if with_call:
                rendezvous = _MultiThreadedRendezvous(state, call, None, deadline)
                return state.response, rendezvous
            else:
                return state.response
        else:
>           raise _InactiveRpcError(state)  # pytype: disable=not-instantiable
E           grpc._channel._InactiveRpcError: <_InactiveRpcError of RPC that terminated with:
E           	status = StatusCode.UNAVAILABLE
E           	details = "failed to connect to all addresses; last error: UNKNOWN: ipv4:84.201.181.26:443: Handshake read failed"
E           	debug_error_string = "UNKNOWN:Error received from peer  {grpc_message:"failed to connect to all addresses; last error: UNKNOWN: ipv4:84.201.181.26:443: Handshake read failed", grpc_status:14, created_time:"2024-05-01T03:00:44.910720305+00:00"}"
E           >

/usr/local/lib/python3.11/site-packages/grpc/_channel.py:1006: _InactiveRpcError

Jed Cunningham 14 hours ago Yep, super easy to reproduce if you shut off wifi on your laptop :slightly_smiling_face: It’s only the 1 test.

tests/providers/yandex/hooks/test_yq.py::TestYandexCloudYqHook::test_select_results

Committer

• [X] I acknowledge that I am a maintainer/committer of the Apache Airflow project.

[dondaum]

Would be happy to take it over.

[Piatachock]

Hi there! Mentioning my fellow colleague from Yandex Query, who made this part of provider @uzhastik Sergei, JFYI

[Taragolis]

I guess this one just required to mock yandexcloud.SDK.client, this one change seems work well (but need to double check)

diff --git a/tests/providers/yandex/hooks/test_yq.py b/tests/providers/yandex/hooks/test_yq.py
index c378c65347..2864642c0e 100644
--- a/tests/providers/yandex/hooks/test_yq.py
+++ b/tests/providers/yandex/hooks/test_yq.py
@@ -31,12 +31,7 @@ IAM_TOKEN = "my_iam_token"
 SERVICE_ACCOUNT_AUTH_KEY_JSON = """{"id":"my_id", "service_account_id":"my_sa1", "private_key":"my_pk"}"""
 
 
-class DummySDK:
-    def __init__(self) -> None:
-        self.client = None
-
-
-class DummyTokenRequester:
+class FakeTokenRequester:
     def get_token(self) -> str:
         return IAM_TOKEN
 
@@ -44,15 +39,12 @@ class DummyTokenRequester:
         return "my_dummy_request"
 
 
-class DummyCreateTokenResponse:
-    iam_token = "zzz"
-
-
 class TestYandexCloudYqHook:
     def _init_hook(self):
         with mock.patch("airflow.hooks.base.BaseHook.get_connection") as mock_get_connection:
             mock_get_connection.return_value = self.connection
-            self.hook = YQHook(default_folder_id="my_folder_id")
+            with mock.patch("airflow.providers.yandex.hooks.yq.yandexcloud.SDK.client"):
+                self.hook = YQHook(default_folder_id="my_folder_id")
 
     def setup_method(self):
         self.connection = Connection(extra={"service_account_json": SERVICE_ACCOUNT_AUTH_KEY_JSON})
@@ -82,7 +74,7 @@ class TestYandexCloudYqHook:
             m.assert_called_once_with("query1")
 
     @responses.activate()
-    @mock.patch("yandexcloud._auth_fabric.get_auth_token_requester", return_value=DummyTokenRequester())
+    @mock.patch("yandexcloud._auth_fabric.get_auth_token_requester", return_value=FakeTokenRequester())
     def test_metadata_token_usage(self, mock_get_auth_token_requester):
         responses.post(
             "https://api.yandex-query.cloud.yandex.net/api/fq/v1/queries",
@@ -101,14 +93,9 @@ class TestYandexCloudYqHook:
         query_id = self.hook.create_query(query_text="select 777", name="my query")
         assert query_id == "query1"
 
-    @mock.patch(
-        "yandex.cloud.iam.v1.iam_token_service_pb2_grpc.IamTokenServiceStub.Create",
-        create=True,
-        new_callable=mock.PropertyMock,
-    )
     @mock.patch("yandexcloud._auth_fabric.__validate_service_account_key")
-    @mock.patch("yandexcloud._auth_fabric.get_auth_token_requester", return_value=DummyTokenRequester())
-    def test_select_results(self, mock_get_auth_token_requester, mock_validate, mock_create_token):
+    @mock.patch("yandexcloud._auth_fabric.get_auth_token_requester", return_value=FakeTokenRequester())
+    def test_select_results(self, mock_get_auth_token_requester, mock_validate):
         with mock.patch.multiple(
             "yandex_query_client.YQHttpClient",
             create_query=mock.DEFAULT,
@@ -120,7 +107,6 @@ class TestYandexCloudYqHook:
         ) as mocks:
             self._init_hook()
             mock_validate.assert_called()
-            mock_create_token.assert_called()
             mock_get_auth_token_requester.assert_called_once_with(
                 service_account_key=json.loads(SERVICE_ACCOUNT_AUTH_KEY_JSON)
             )

cc: @dondaum @uzhastik @Piatachock