code: def test_athena_query_wait(sagemaker_session_mock, query):
      query._current_query_execution_id = "query_id"
      query.wait()
      sagemaker_session_mock.wait_for_athena_query.assert_called_with(
          query_execution_id="query_id"
      )
  
imports: from __future__ import absolute_import
  import pandas
  import pytest
  from mock import Mock, patch, MagicMock
  from botocore.exceptions import ProfileNotFound
  from sagemaker.feature_store.feature_definition import FractionalFeatureDefinition, IntegralFeatureDefinition, StringFeatureDefinition, FeatureTypeEnum, VectorCollectionType, SetCollectionType, ListCollectionType
  from sagemaker.feature_store.feature_group import FeatureGroup, IngestionManagerPandas, AthenaQuery, IngestionError
  from sagemaker.feature_store.inputs import FeatureParameter, DeletionModeEnum, TtlDuration, OnlineStoreConfigUpdate, OnlineStoreStorageTypeEnum
  from tests.unit import SAGEMAKER_CONFIG_FEATURE_GROUP
label_desc: Use image node_id 44 for calling a global function with example usage: test_athena_query_wait(sagemaker_session_mock, query) without return types