code: def test_databricks_uc_datasource():
      MockResponse = namedtuple(
          "Response", "raise_for_status json content"
      )
  
      MockChunk = namedtuple(
          "Chunk", "index, row_count byte_count data"
      )
  
      token = "test_token"
      warehouse_id = "test_warehouse_id"
      catalog = "catalog1"
      schema = "db1"
      query = "select * from table1"
      expected_result_df = pd.DataFrame(
          {
              "c1": range(10000),
              "c2": map(lambda x: "str" + str(x), range(10000)),
          }
      )
      rows_per_chunk = 700
  
      @contextmanager
      def setup_mock():
          mock_chunks = []
  
          num_rows = len(expected_result_df)
          cur_pos = 0
          index = 0
  
          while cur_pos < num_rows:
              if cur_pos + rows_per_chunk <= num_rows:
                  chunk_rows = rows_per_chunk
              else:
                  chunk_rows = num_rows - cur_pos
  
              chunk_df = expected_result_df[
                  cur_pos : (cur_pos + chunk_rows)
              ]
              chunk_pa_table = pa.Table.from_pandas(chunk_df)
              sink = pa.BufferOutputStream()
              with pa.ipc.new_stream(
                  sink, chunk_pa_table.schema
              ) as writer:
                  writer.write_table(chunk_pa_table)
  
              chunk_data = sink.getvalue()
  
              mock_chunks.append(
                  MockChunk(
                      index=index,
                      row_count=chunk_rows,
                      byte_count=len(chunk_data),
                      data=chunk_data,
                  )
              )
              index += 1
              cur_pos += rows_per_chunk
  
          chunk_meta_json = [
              {
                  "chunk_index": index,
                  "row_count": mock_chunk.row_count,
                  "byte_count": mock_chunk.byte_count,
              }
              for index, mock_chunk in enumerate(mock_chunks)
          ]
          chunk_meta_json.reverse()
          valid_statement_ids = set()
  
          def request_post_mock(url, data=None, json=None, **kwargs):
              import json as jsonlib
  
              auth = kwargs["auth"]
              headers = kwargs["headers"]
  
              if url == "https://test_shard/api/2.0/sql/statements/":
                  assert auth == ("token", token)
                  assert headers == {"Content-Type": "application/json"}
                  assert jsonlib.loads(data) == {
                      "statement": query,
                      "warehouse_id": warehouse_id,
                      "wait_timeout": "0s",
                      "disposition": "EXTERNAL_LINKS",
                      "format": "ARROW_STREAM",
                      "catalog": catalog,
                      "schema": schema,
                  }
  
                  statement_id = uuid.uuid4().hex
                  valid_statement_ids.add(statement_id)
  
                  return MockResponse(
                      raise_for_status=lambda: None,
                      json=lambda: {
                          "statement_id": statement_id,
                          "status": {"state": "PENDING"},
                      },
                      content=b"",
                  )
  
              assert False, "Invalid request."
  
          def request_get_mock(url, params=None, **kwargs):
              auth = kwargs["auth"]
              headers = kwargs["headers"]
  
              if match := re.match(
                  r"^https://test_shard/api/2\.0/sql/statements/([^/]*)/$",
                  url,
              ):
                  statement_id = match.group(1)
                  assert auth == ("token", token)
                  assert headers == {"Content-Type": "application/json"}
  
                  assert statement_id in valid_statement_ids
  
                  return MockResponse(
                      raise_for_status=lambda: None,
                      json=lambda: {
                          "status": {"state": "SUCCEEDED"},
                          "manifest": {
                              "truncated": False,
                              "chunks": chunk_meta_json,
                          },
                      },
                      content=None,
                  )
  
              if match := re.match(
                  r"^https://test_shard/api/2\.0/sql/"
                  r"statements/([^/]*)/result/chunks/([^/]*)$",
                  url,
              ):
                  assert auth == ("token", token)
                  assert headers == {"Content-Type": "application/json"}
  
                  chunk_index = match.group(2)
  
                  external_link = (
                      f"https://test_external_link/{chunk_index}"
                  )
                  return MockResponse(
                      raise_for_status=lambda: None,
                      json=lambda: {
                          "external_links": [
                              {
                                  "external_link": external_link,
                              }
                          ]
                      },
                      content=None,
                  )
  
              if match := re.match(
                  r"^https://test_external_link/([^/]*)$", url
              ):
                  assert auth is None
                  assert headers is None
  
                  chunk_index = int(match.group(1))
  
                  return MockResponse(
                      raise_for_status=lambda: None,
                      json=lambda: None,
                      content=mock_chunks[chunk_index].data,
                  )
  
              assert False, "Invalid request."
  
          with mock.patch("requests.get", request_get_mock), mock.patch(
              "requests.post", request_post_mock
          ), mock.patch.dict(
              os.environ,
              {
                  "DATABRICKS_HOST": "test_shard",
                  "DATABRICKS_TOKEN": token,
              },
          ):
              yield
  
      setup_mock_fn_path = os.path.join(
          tempfile.mkdtemp(), "setup_mock_fn.pkl"
      )
      with open(setup_mock_fn_path, "wb") as fp:
          pickle.dump(setup_mock, fp)
  
      MOCK_ENV = (
          "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
      )
      with setup_mock(), mock.patch.dict(
          os.environ, {MOCK_ENV: setup_mock_fn_path}
      ):
          # shut down existing Ray local cluster and
          # recreate Ray local cluster so that Ray head node could get the
          # "MOCK_ENV"
          ray.shutdown()
          ray.init()
  
          # test query with a table name
          result = ray.data.read_databricks_tables(
              warehouse_id=warehouse_id,
              table="table1",
              catalog="catalog1",
              schema="db1",
              parallelism=5,
          ).to_pandas()
  
          pd.testing.assert_frame_equal(result, expected_result_df)
  
          # test query with SQL
          result = ray.data.read_databricks_tables(
              warehouse_id=warehouse_id,
              query="select * from table1",
              catalog="catalog1",
              schema="db1",
              parallelism=5,
          ).to_pandas()
  
          pd.testing.assert_frame_equal(result, expected_result_df)
  
          # test larger parallelism
          result = ray.data.read_databricks_tables(
              warehouse_id=warehouse_id,
              query="select * from table1",
              catalog="catalog1",
              schema="db1",
              parallelism=100,
          ).to_pandas()
  
          pd.testing.assert_frame_equal(result, expected_result_df)
  
imports: import os
  import re
  import sqlite3
  import tempfile
  import uuid
  from collections import namedtuple
  from contextlib import contextmanager
  from typing import Generator
  from unittest import mock
  import pandas
  import pyarrow
  import pytest
  import ray
  import ray.cloudpickle
label_desc: Use image node_id 6 for calling a global function with example usage: test_databricks_uc_datasource() and returns: MockResponse, MockResponse, MockResponse, MockResponse