-
Notifications
You must be signed in to change notification settings - Fork 106
Expand file tree
/
Copy pathshared_test_pipeline.py
More file actions
91 lines (78 loc) · 2.72 KB
/
shared_test_pipeline.py
File metadata and controls
91 lines (78 loc) · 2.72 KB
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
import unittest
from feldera.runtime_config import RuntimeConfig
from feldera.testutils import (
unique_pipeline_name,
FELDERA_TEST_NUM_WORKERS,
FELDERA_TEST_NUM_HOSTS,
)
from tests import TEST_CLIENT
from feldera import PipelineBuilder, Pipeline
def sql(text_or_iterable):
"""
Decorator to attach SQL (string or list/tuple of strings) to a test method.
"""
def _wrap(fn):
fn.SQL = text_or_iterable
return fn
return _wrap
class SharedTestPipeline(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls._ddls = []
cls.client = TEST_CLIENT
cls.pipeline_name = unique_pipeline_name(cls.__name__)
for attr in dir(cls):
if not attr.startswith("test_"):
continue
func = getattr(cls, attr)
# Check for enterprise_only decorator
is_enterprise_only = getattr(func, "_enterprise_only", False)
if (
is_enterprise_only
and not cls.client.get_config().edition.is_enterprise()
):
continue # Skip DDL for enterprise-only tests if not enterprise
ddl = getattr(func, "SQL", getattr(func, "__doc__", None))
if ddl and ddl.strip() not in cls._ddls:
cls._ddls.append(ddl.strip())
if not hasattr(cls, "_pipeline"):
cls.ddl = "\n".join(cls._ddls)
cls._pipeline = PipelineBuilder(
cls.client,
cls.pipeline_name,
cls.ddl,
runtime_config=RuntimeConfig(
workers=FELDERA_TEST_NUM_WORKERS,
hosts=FELDERA_TEST_NUM_HOSTS,
logging="debug",
),
).create_or_replace()
def setUp(self):
p = PipelineBuilder(
self.client,
unique_pipeline_name(self._testMethodName),
sql=self.ddl,
runtime_config=RuntimeConfig(
workers=FELDERA_TEST_NUM_WORKERS,
hosts=FELDERA_TEST_NUM_HOSTS,
logging="debug",
),
).create_or_replace()
self.p = p
def tearDown(self):
self.p.stop(force=True)
self.p.clear_storage()
@property
def pipeline(self) -> Pipeline:
return self.p
def new_pipeline_with_suffix(self, suffix: str) -> Pipeline:
return PipelineBuilder(
self.client,
unique_pipeline_name(f"{self._testMethodName}_{suffix}"),
sql=self.ddl,
runtime_config=RuntimeConfig(
workers=FELDERA_TEST_NUM_WORKERS,
hosts=FELDERA_TEST_NUM_HOSTS,
logging="debug",
),
).create_or_replace()