code: def test_generate_ingestion_flow_for_athena_input():
      athena_df = AthenaDatasetDefinition(
          catalog="AwsDataCatalog",
          database="sagemaker_db",
          query_string='SELECT * FROM "table"',
          output_s3_uri="s3://my-bucket/athena/",
          output_format="parquet",
      )
  
      (
          flow,
          output_name,
      ) = generate_data_ingestion_flow_from_athena_dataset_definition(
          "test-athena",
          athena_dataset_definition=athena_df,
          operator_version="0.2",
          schema={"k1": "string", "k2": "long"},
      )
  
      assert flow["metadata"] == {"version": 1, "disable_limits": False}
      assert len(flow["nodes"]) == 2
      assert output_name.endswith(
          flow["nodes"][1]["outputs"][0]["name"]
      )
      # reset node id
      flow["nodes"][0]["node_id"] = "123456"
      flow["nodes"][1]["node_id"] = "7891011"
      flow["nodes"][1]["inputs"][0]["node_id"] = "123456"
      assert flow["nodes"] == [
          {
              "node_id": "123456",
              "type": "SOURCE",
              "operator": "sagemaker.athena_source_0.2",
              "parameters": {
                  "dataset_definition": {
                      "datasetSourceType": "Athena",
                      "name": "test-athena",
                      "catalogName": athena_df.catalog,
                      "databaseName": athena_df.database,
                      "queryString": athena_df.query_string,
                      "s3OutputLocation": athena_df.output_s3_uri,
                      "outputFormat": athena_df.output_format,
                  }
              },
              "inputs": [],
              "outputs": [{"name": "default"}],
          },
          {
              "node_id": "7891011",
              "type": "TRANSFORM",
              "operator": "sagemaker.spark.infer_and_cast_type_0.2",
              "trained_parameters": {"k1": "string", "k2": "long"},
              "parameters": {},
              "inputs": [
                  {
                      "name": "default",
                      "node_id": "123456",
                      "output_name": "default",
                  }
              ],
              "outputs": [{"name": "default"}],
          },
      ]
  
imports: from __future__ import absolute_import
  from sagemaker.wrangler.ingestion import generate_data_ingestion_flow_from_s3_input, generate_data_ingestion_flow_from_athena_dataset_definition, generate_data_ingestion_flow_from_redshift_dataset_definition
  from sagemaker.dataset_definition.inputs import RedshiftDatasetDefinition, AthenaDatasetDefinition
label_desc: Use image node_id 2 for calling a global function with example usage: test_generate_ingestion_flow_for_athena_input() without return types