hangshuo652
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
from __future__ import annotations
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
|
|
@dataclass
|
|
class SparkConfig:
|
|
num_records: int = 100
|
|
replication: str = "key_varied"
|
|
key_field: str = ""
|
|
edge_cases: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class TestCase:
|
|
id: str
|
|
fields: dict = field(default_factory=dict)
|
|
coverage_targets: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class TestSuite:
|
|
schema: Optional[dict] = None
|
|
test_cases: list[TestCase] = field(default_factory=list)
|
|
spark_config: Optional[SparkConfig] = None
|
|
|
|
def with_spark(self, num_records: int = 1000, key_field: str = "") -> "TestSuite":
|
|
self.spark_config = SparkConfig(num_records=num_records, key_field=key_field)
|
|
return self
|
|
|
|
@property
|
|
def has_spark(self) -> bool:
|
|
return self.spark_config is not None
|
|
|
|
|
|
_tc = TestCase(id="TC-001", fields={"BR-AMT": 1500000, "BR-TYPE": "A"})
|
|
assert _tc.id == "TC-001"
|
|
assert _tc.fields["BR-AMT"] == 1500000
|
|
|
|
_ts = TestSuite(test_cases=[_tc], spark_config=SparkConfig(num_records=1000))
|
|
assert _ts.spark_config.num_records == 1000
|