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