datacloud-customcode-python-sdk/tests/spark/test_session_provider.py at 69f262d3a3931e51856bcbb5afff192edcdf23e8 · forcedotcom/datacloud-customcode-python-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
from __future__ import annotations

from typing import TYPE_CHECKING, Any

import builtins

from datacustomcode.client import Client
from datacustomcode.config import (
    AccessLayerObjectConfig,
    ClientConfig,
    SparkConfig,
    SparkProviderConfig,
)
from datacustomcode.io.reader.base import BaseDataCloudReader
from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
from datacustomcode.spark.base import BaseSparkSessionProvider

if TYPE_CHECKING:
    from pyspark.sql import DataFrame as PySparkDataFrame


class _Sentinel:
    pass


SENTINEL_SPARK = _Sentinel()


class MockReader(BaseDataCloudReader):
    CONFIG_NAME = "MockReader"
    last_spark: Any | None = None

    def __init__(self, spark):
        super().__init__(spark)
        MockReader.last_spark = spark

    def read_dlo(self, name: str):  # type: ignore[override]
        raise NotImplementedError

    def read_dmo(self, name: str):  # type: ignore[override]
        raise NotImplementedError


class MockWriter(BaseDataCloudWriter):
    CONFIG_NAME = "MockWriter"
    last_spark: Any | None = None

    def __init__(self, spark):
        super().__init__(spark)
        MockWriter.last_spark = spark

    def write_to_dlo(
        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
    ) -> None:  # type: ignore[override]
        raise NotImplementedError

    def write_to_dmo(
        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
    ) -> None:  # type: ignore[override]
        raise NotImplementedError


class FakeProvider(BaseSparkSessionProvider):
    CONFIG_NAME = "FakeProvider"

    def get_session(self, spark_config: SparkConfig):  # type: ignore[override]
        return SENTINEL_SPARK


def _reset_singleton():
    # Reset Client singleton between tests
    Client._instance = None  # type: ignore[attr-defined]


def test_client_uses_provider_from_config(monkeypatch):
    _reset_singleton()

    cfg = ClientConfig(
        reader_config=AccessLayerObjectConfig(
            type_config_name=MockReader.CONFIG_NAME, options={}
        ),
        writer_config=AccessLayerObjectConfig(
            type_config_name=MockWriter.CONFIG_NAME, options={}
        ),
        spark_config=SparkConfig(app_name="test-app", master=None, options={}),
        spark_provider_config=SparkProviderConfig(
            type_config_name=FakeProvider.CONFIG_NAME, options={}
        ),
    )

    from datacustomcode.config import config as global_config

    global_config.update(cfg)

    Client()  # noqa: F841
    assert MockReader.last_spark is SENTINEL_SPARK
    assert MockWriter.last_spark is SENTINEL_SPARK

class ExplicitProvider(BaseSparkSessionProvider):
    CONFIG_NAME = "ExplicitProvider"

    def get_session(self, spark_config: SparkConfig):  # type: ignore[override]
        return SENTINEL_SPARK


def test_client_explicit_provider_overrides_config(monkeypatch):
    _reset_singleton()

    cfg = ClientConfig(
        reader_config=AccessLayerObjectConfig(
            type_config_name=MockReader.CONFIG_NAME, options={}
        ),
        writer_config=AccessLayerObjectConfig(
            type_config_name=MockWriter.CONFIG_NAME, options={}
        ),
        spark_config=SparkConfig(app_name="test-app", master=None, options={}),
        spark_provider_config=None,
    )

    from datacustomcode.config import config as global_config

    global_config.update(cfg)

    provider = ExplicitProvider()
    Client(spark_provider=provider)  # noqa: F841
    assert MockReader.last_spark is SENTINEL_SPARK
    assert MockWriter.last_spark is SENTINEL_SPARK