Repository: ray-project/ray
File: python/ray/data/tests/datasource/test_databricks_uc_datasource.py (Python)

Code before:
                schema=test_data["schema"],
                override_num_blocks=5,
            ).to_pandas()

            assert rows_same(result, test_data["expected_df"])

    @pytest.mark.parametrize("num_blocks", [5, 100])
    def test_read_with_different_parallelism(self, test_data, num_blocks):
        """Test reading data with different parallelism settings."""
        with self._setup_integration_test(test_data):
            result = ray.data.read_databricks_tables(
                warehouse_id=test_data["warehouse_id"],
                query=test_data["query"],
                catalog=test_data["catalog"],
                schema=test_data["schema"],
                override_num_blocks=num_blocks,
            ).to_pandas()

            assert rows_same(result, test_data["expected_df"])


class TestDatabricksUCDatasourceCredentials:
    """Tests for credential provider handling."""

    def test_schema_name_does_not_shadow_datasource_fields(self, requests_mocker):
        """Test that schema name is stored without using the `schema` attribute."""
        requests_mocker["post"].return_value = mock.Mock(
            status_code=200,
            raise_for_status=lambda: None,
            json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
        )
        requests_mocker["get"].return_value = mock.Mock(
            status_code=200,
            raise_for_status=lambda: None,
            json=lambda: {
                "status": {"state": "SUCCEEDED"},
                "manifest": {"truncated": False, "chunks": []},
            },
        )

        provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
        datasource = DatabricksUCDatasource(
            warehouse_id="test_warehouse",
            catalog="test_catalog",
            schema="test_schema",
            query="SELECT 1",
            credential_provider=provider,
        )

        assert datasource.schema_name == "test_schema"
        assert "schema" not in datasource.__dict__

Review comment (suggestion):
```suggestion
        """Test that schema name is stored without using the `schema` attribute.
        
        This is a regression test for https://github.com/ray-project/ray/issues/46481.
        """
```

Code after:
                schema=test_data["schema"],
                override_num_blocks=5,
            ).to_pandas()

            assert rows_same(result, test_data["expected_df"])

    @pytest.mark.parametrize("num_blocks", [5, 100])
    def test_read_with_different_parallelism(self, test_data, num_blocks):
        """Test reading data with different parallelism settings."""
        with self._setup_integration_test(test_data):
            result = ray.data.read_databricks_tables(
                warehouse_id=test_data["warehouse_id"],
                query=test_data["query"],
                catalog=test_data["catalog"],
                schema=test_data["schema"],
                override_num_blocks=num_blocks,
            ).to_pandas()

            assert rows_same(result, test_data["expected_df"])


class TestDatabricksUCDatasourceCredentials:
    """Tests for credential provider handling."""

    def test_schema_name_does_not_shadow_datasource_fields(self, requests_mocker):
        """Test that schema name is stored without using the `schema` attribute.
        This is a regression test for https://github.com/ray-project/ray/issues/46481.
        """
        requests_mocker["post"].return_value = mock.Mock(
            status_code=200,
            raise_for_status=lambda: None,
            json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
        )
        requests_mocker["get"].return_value = mock.Mock(
            status_code=200,
            raise_for_status=lambda: None,
            json=lambda: {
                "status": {"state": "SUCCEEDED"},
                "manifest": {"truncated": False, "chunks": []},
            },
        )

        provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
        datasource = DatabricksUCDatasource(
            warehouse_id="test_warehouse",
            catalog="test_catalog",
            schema="test_schema",
            query="SELECT 1",
            credential_provider=provider,
        )