diff --git a/.commandcode/settings.json b/.commandcode/settings.json new file mode 100644 index 0000000..4a114a4 --- /dev/null +++ b/.commandcode/settings.json @@ -0,0 +1,9 @@ +{ + "permissions": { + "allow": [ + "Shell(python -m ruff check src tests 2 >& 1)" + ], + "deny": [], + "defaultMode": "default" + } +} \ No newline at end of file diff --git a/.commandcode/taste/coding/taste.md b/.commandcode/taste/coding/taste.md new file mode 100644 index 0000000..145a430 --- /dev/null +++ b/.commandcode/taste/coding/taste.md @@ -0,0 +1,4 @@ +# Coding + +- For Python projects lacking lint/type tooling, endorsed adding ruff (lint + import sorting) plus mypy with a lenient starter config as part of the work. Confidence: 0.7 +- Prefers behavior-preserving refactors: deduplicate and split overly complex functions without changing the public API surface, module paths, or documented architecture (keeps docs/roadmaps referencing the layout valid). Confidence: 0.6 diff --git a/.commandcode/taste/taste.md b/.commandcode/taste/taste.md new file mode 100644 index 0000000..23afde4 --- /dev/null +++ b/.commandcode/taste/taste.md @@ -0,0 +1,4 @@ +# Taste + +## Communication +- After reviewing a detailed plan, approves with a terse "go ahead" and delegates full autonomous execution; wants key decisions surfaced as a small set of upfront questions before implementation, not repeated check-ins mid-work. Confidence: 0.6 diff --git a/.commandcode/taste/workflow/taste.md b/.commandcode/taste/workflow/taste.md new file mode 100644 index 0000000..953b3be --- /dev/null +++ b/.commandcode/taste/workflow/taste.md @@ -0,0 +1,8 @@ +# Workflow + +- For refactoring tasks, capture a baseline test + coverage run before changing anything, re-run the full test suite after each phase so any breakage is attributable to a single phase, and run the suite twice at the end to catch test-order coupling. Confidence: 0.6 +- Track multi-phase work as an explicit phased todo list (tooling → mechanical cleanup → dedupe → complexity → architecture → verification) rather than one undifferentiated task. Confidence: 0.5 +- Verify the local copy is current (git log/status) before deep codebase analysis or planning; if it's stale, re-run exploration from scratch — explicitly ignoring prior findings — and update the plan, instead of incrementally patching a plan built on outdated code. Confidence: 0.6 +- When a long-running command (e.g., a test suite) appears stuck, don't block on it — explicitly asked to "proceed with next step"; keep making forward progress, run the stuck operation in the background and/or with a per-test timeout so hangs get diagnosed without halting the work. Confidence: 0.7 +- When the local environment can't run the full test suite (e.g., Windows dev box missing bash/WSL or native deps), don't stall on fixing it locally — explicitly directed to "proceed with all the changes for now. I will test the same inside my Github CI instead." Rely on static checks (ruff, mypy) and runnable local subsets meanwhile, and treat GitHub CI as the arbiter for full-suite validation. Confidence: 0.8 +- Local dev machine is Windows with known environmental test failures (tests spawning a POSIX shell via `bash -c` with no WSL installed, native Spark writers needing Hadoop winutils); classify those as pre-existing/environmental rather than regressions, and diff against an expected-failure baseline instead of chasing a fully green local run. Confidence: 0.7 diff --git a/pyproject.toml b/pyproject.toml index 9a9f052..cc072cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" name = "testbricks" description = "A set of proxy objects to facilitate testing of Databricks notebooks in CI/CD pipelines" dynamic = ["version"] -requires-python = ">=3.8" +requires-python = ">=3.10" authors = [ {name = "Karan Gupta", email = "gkaran184@gmail.com"}, ] @@ -21,8 +21,6 @@ classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -30,7 +28,28 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] +[project.optional-dependencies] +dev = ["pytest", "coverage", "ruff", "mypy"] + [tool.setuptools.packages.find] where = ["src"] [tool.setuptools_scm] + +[tool.ruff] +line-length = 100 +target-version = "py310" +src = ["src", "tests"] + +[tool.ruff.lint] +select = ["E", "W", "F", "I", "B"] + +[tool.ruff.lint.per-file-ignores] +# Databricks notebook fixtures: spark/dbutils are injected at runtime. +"tests/e2e_workflow/notebooks/*.py" = ["F821"] + +[tool.mypy] +# numpy/pandas stubs use Python 3.12 syntax, so the checker must run as 3.12 +# even though the package supports 3.10+. +python_version = "3.12" +ignore_missing_imports = true diff --git a/requirements.txt b/requirements.txt index 2cbe210..118b009 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,6 @@ pyarrow numpy py4j coverage +ruff +mypy +pytest-timeout diff --git a/src/testbricks/catalog/__init__.py b/src/testbricks/catalog/__init__.py index 293c7ac..c44bd15 100644 --- a/src/testbricks/catalog/__init__.py +++ b/src/testbricks/catalog/__init__.py @@ -1,3 +1,10 @@ +from .csv_options import ( + CSV_OPTION_KEYS, + DEFAULT_CSV_READ_OPTIONS, + normalize_csv_options, + option_flag, + option_lookup, +) from .errors import InvalidTableNameError, SchemaMismatchError, SparkProxyError from .identifier import TableIdentifier from .spark_catalog import CatalogFacade @@ -5,12 +12,17 @@ from .table_catalog import TableCatalog __all__ = [ + "CSV_OPTION_KEYS", "CatalogFacade", + "DEFAULT_CSV_READ_OPTIONS", "InvalidTableNameError", "SchemaMismatchError", "SparkProxyError", "TableCatalog", "TableIdentifier", "is_maintenance_noop", + "normalize_csv_options", + "option_flag", + "option_lookup", "rewrite_from_join_identifiers", ] diff --git a/src/testbricks/catalog/csv_options.py b/src/testbricks/catalog/csv_options.py new file mode 100644 index 0000000..f68b991 --- /dev/null +++ b/src/testbricks/catalog/csv_options.py @@ -0,0 +1,69 @@ +"""Canonical CSV option keys, aliases, and case-insensitive lookups. + +This module is the single source of truth for which ``option(...)`` keys the +CSV catalog honors, how Spark-style aliases map to canonical names, and how +options are matched case-insensitively. +""" + +from __future__ import annotations + +from typing import Mapping, Optional + +DEFAULT_CSV_READ_OPTIONS: dict[str, str] = {"header": "true", "inferSchema": "true"} + +CSV_OPTION_KEYS = frozenset( + { + "delimiter", + "sep", + "quote", + "escape", + "nullvalue", + "dateformat", + "timestampformat", + "header", + } +) + +_OPTION_ALIASES = { + "sep": "delimiter", + "delimiter": "delimiter", + "quote": "quote", + "escape": "escape", + "nullvalue": "nullValue", + "dateformat": "dateFormat", + "timestampformat": "timestampFormat", + "header": "header", +} + + +def normalize_csv_options(options: Optional[Mapping[str, str]]) -> dict[str, str]: + """Keep only known CSV options, mapping aliases to canonical names.""" + if not options: + return {} + normalized: dict[str, str] = {} + for key, value in options.items(): + canonical = _OPTION_ALIASES.get(key.lower()) + if canonical is None: + continue + normalized[canonical] = str(value) + return normalized + + +def option_lookup(options: Optional[Mapping[str, str]], *names: str) -> Optional[str]: + """Case-insensitive lookup returning the first match as ``str``.""" + if not options: + return None + lowered = {key.lower(): value for key, value in options.items()} + for name in names: + if name.lower() in lowered: + return str(lowered[name.lower()]) + return None + + +def option_flag(options: Mapping[str, str], *names: str) -> bool: + """Case-insensitive truthy check for Spark-style boolean option values.""" + wanted = {name.lower() for name in names} + for key, value in options.items(): + if key.lower() in wanted: + return str(value).lower() in {"true", "1", "yes"} + return False diff --git a/src/testbricks/catalog/errors.py b/src/testbricks/catalog/errors.py index 02db9ea..7e4998e 100644 --- a/src/testbricks/catalog/errors.py +++ b/src/testbricks/catalog/errors.py @@ -1,7 +1,15 @@ -"""Errors raised by the SparkProxy CSV catalog.""" +"""Errors raised by the CSV catalog layer. +All catalog errors derive from both ``TestbricksError`` and the hierarchy +below. ``InvalidTableNameError`` and ``SchemaMismatchError`` additionally +inherit ``ValueError`` because callers legitimately treat invalid table +names and schema mismatches as value-level input errors. +""" -class SparkProxyError(Exception): +from testbricks.errors import TestbricksError + + +class SparkProxyError(TestbricksError): """Base error for SparkProxy catalog and table operations.""" diff --git a/src/testbricks/catalog/identifier.py b/src/testbricks/catalog/identifier.py index 5967be4..ae45ab1 100644 --- a/src/testbricks/catalog/identifier.py +++ b/src/testbricks/catalog/identifier.py @@ -7,11 +7,16 @@ from .errors import InvalidTableNameError +def strip_wrappers(text: str, wrapper_chars: str = "`") -> str: + """Strip one layer of symmetric wrapper characters (e.g. backticks, quotes).""" + text = text.strip() + if len(text) >= 2 and text[0] == text[-1] and text[0] in wrapper_chars: + return text[1:-1] + return text + + def _strip_identifier_part(part: str) -> str: - part = part.strip() - if len(part) >= 2 and part[0] == "`" and part[-1] == "`": - return part[1:-1] - return part + return strip_wrappers(part) def _split_qualified_name(table_name: str) -> list[str]: diff --git a/src/testbricks/catalog/spark_catalog.py b/src/testbricks/catalog/spark_catalog.py index d77b02a..65bf65d 100644 --- a/src/testbricks/catalog/spark_catalog.py +++ b/src/testbricks/catalog/spark_catalog.py @@ -22,6 +22,9 @@ class CatalogFacade: def __init__(self, tables: TableCatalog): self._tables = tables + def exists(self, ident: TableIdentifier) -> bool: + return self._tables.exists(ident) + def tableExists(self, tableName: str, dbName: Optional[str] = None) -> bool: try: ident = ( @@ -67,6 +70,3 @@ def listDatabases(self, pattern: Optional[str] = None) -> List[Database]: for name in self._tables.iter_schema_names() if _matches(name, pattern) ] - - def __getattr__(self, name): - return getattr(self._tables, name) diff --git a/src/testbricks/catalog/table_catalog.py b/src/testbricks/catalog/table_catalog.py index 04963fa..28eba85 100644 --- a/src/testbricks/catalog/table_catalog.py +++ b/src/testbricks/catalog/table_catalog.py @@ -13,6 +13,7 @@ from pyspark.sql import DataFrame, SparkSession from pyspark.sql.utils import AnalysisException +from .csv_options import DEFAULT_CSV_READ_OPTIONS, normalize_csv_options, option_lookup from .errors import SchemaMismatchError from .identifier import TableIdentifier @@ -21,20 +22,6 @@ _OVERWRITE_MODES = frozenset({"overwrite"}) _IGNORE_MODES = frozenset({"ignore"}) -_DEFAULT_READ_OPTIONS = {"header": "true", "inferSchema": "true"} -_CSV_OPTION_KEYS = frozenset( - { - "delimiter", - "sep", - "quote", - "escape", - "nullvalue", - "dateformat", - "timestampformat", - "header", - } -) - class TableCatalog: """Maps ``schema.table`` identifiers to local CSV files and Spark temp views.""" @@ -94,7 +81,7 @@ def read_csv( options: Optional[Mapping[str, str]] = None, ) -> DataFrame: merged = { - **_DEFAULT_READ_OPTIONS, + **DEFAULT_CSV_READ_OPTIONS, **self.csv_options_for(ident), **dict(options or {}), } @@ -128,7 +115,7 @@ def save_dataframe( return stored = self.csv_options_for(ident) if exists else {} - incoming = _normalize_csv_options(csv_options) + incoming = normalize_csv_options(csv_options) if save_mode in _APPEND_MODES and exists: effective_options = {**stored, **incoming} else: @@ -138,53 +125,25 @@ def save_dataframe( else: effective_options["header"] = "false" - new_pdf = dataframe.toPandas() - new_pdf = _format_temporal_columns(new_pdf, effective_options) + new_pdf = _format_temporal_columns(dataframe.toPandas(), effective_options) if replace_where: if save_mode != "overwrite": raise ValueError( - "option('replaceWhere') requires mode('overwrite'); " - f"got mode '{mode}'." + f"option('replaceWhere') requires mode('overwrite'); got mode '{mode}'." ) if exists: - existing_pdf = _pandas_read_csv(csv_path, stored or effective_options) - remaining = _apply_replace_where(existing_pdf, replace_where) - aligned = _align_columns_for_concat(remaining, new_pdf) - new_pdf = pd.concat(aligned, ignore_index=True) - elif exists and save_mode == "overwrite": - existing_pdf = _pandas_read_csv(csv_path, stored or effective_options) - if _schema_incompatible(existing_pdf, new_pdf) and not overwrite_schema: - raise SchemaMismatchError( - f"Cannot overwrite '{ident}' with an incompatible schema unless " - "overwriteSchema=true. " - f"Existing columns={list(existing_pdf.columns)}, " - f"new columns={list(new_pdf.columns)}" + new_pdf = self._apply_replace_where( + csv_path, effective_options, new_pdf, replace_where ) + elif exists and save_mode == "overwrite": + self._check_overwrite_schema( + ident, csv_path, effective_options, new_pdf, overwrite_schema + ) elif save_mode in _APPEND_MODES and exists: - existing_pdf = _pandas_read_csv(csv_path, stored or effective_options) - type_conflict = _overlapping_type_conflicts(existing_pdf, new_pdf) - column_mismatch = set(existing_pdf.columns) != set(new_pdf.columns) - if type_conflict: - raise SchemaMismatchError( - f"Cannot append to '{ident}': schema mismatch. " - f"Existing columns={list(existing_pdf.columns)}, " - f"new columns={list(new_pdf.columns)}" - ) - if column_mismatch and merge_schema: - aligned = _align_columns_for_concat(existing_pdf, new_pdf) - new_pdf = pd.concat(aligned, ignore_index=True) - elif column_mismatch: - raise SchemaMismatchError( - f"Cannot append to '{ident}': schema mismatch. " - f"Existing columns={list(existing_pdf.columns)}, " - f"new columns={list(new_pdf.columns)}" - ) - else: - new_pdf = pd.concat( - [existing_pdf, new_pdf[existing_pdf.columns]], - ignore_index=True, - ) + new_pdf = self._append_with_schema_checks( + ident, csv_path, effective_options, new_pdf, merge_schema + ) self._write_csv_atomic(new_pdf, csv_path, header=header, options=effective_options) self._persist_csv_options(ident, effective_options) @@ -192,6 +151,64 @@ def save_dataframe( ident.view_name ) + def _apply_replace_where( + self, + csv_path: str, + read_options: Mapping[str, str], + new_pdf: pd.DataFrame, + replace_where: str, + ) -> pd.DataFrame: + existing_pdf = _pandas_read_csv(csv_path, read_options) + remaining = _drop_rows_matching(existing_pdf, replace_where) + aligned = _align_columns_for_concat(remaining, new_pdf) + return pd.concat(aligned, ignore_index=True) + + def _check_overwrite_schema( + self, + ident: TableIdentifier, + csv_path: str, + read_options: Mapping[str, str], + new_pdf: pd.DataFrame, + overwrite_schema: bool, + ) -> None: + existing_pdf = _pandas_read_csv(csv_path, read_options) + if _schema_incompatible(existing_pdf, new_pdf) and not overwrite_schema: + raise SchemaMismatchError( + f"Cannot overwrite '{ident}' with an incompatible schema unless " + "overwriteSchema=true. " + f"Existing columns={list(existing_pdf.columns)}, " + f"new columns={list(new_pdf.columns)}" + ) + + def _append_with_schema_checks( + self, + ident: TableIdentifier, + csv_path: str, + read_options: Mapping[str, str], + new_pdf: pd.DataFrame, + merge_schema: bool, + ) -> pd.DataFrame: + existing_pdf = _pandas_read_csv(csv_path, read_options) + mismatch_details = ( + f"Existing columns={list(existing_pdf.columns)}, new columns={list(new_pdf.columns)}" + ) + if _overlapping_type_conflicts(existing_pdf, new_pdf): + raise SchemaMismatchError( + f"Cannot append to '{ident}': schema mismatch. {mismatch_details}" + ) + column_mismatch = set(existing_pdf.columns) != set(new_pdf.columns) + if column_mismatch and merge_schema: + aligned = _align_columns_for_concat(existing_pdf, new_pdf) + return pd.concat(aligned, ignore_index=True) + if column_mismatch: + raise SchemaMismatchError( + f"Cannot append to '{ident}': schema mismatch. {mismatch_details}" + ) + return pd.concat( + [existing_pdf, new_pdf[existing_pdf.columns]], + ignore_index=True, + ) + def _persist_csv_options(self, ident: TableIdentifier, options: Mapping[str, str]) -> None: payload = {key: str(value) for key, value in options.items()} self._csv_options[str(ident)] = dict(payload) @@ -232,59 +249,26 @@ def _normalize_save_mode(mode: Optional[str]) -> str: if normalized in _IGNORE_MODES: return "ignore" raise ValueError( - f"Unknown save mode '{mode}'. Expected overwrite, append, ignore, " - "error, or errorIfExists." + f"Unknown save mode '{mode}'. Expected overwrite, append, ignore, error, or errorIfExists." ) -def _normalize_csv_options(options: Optional[Mapping[str, str]]) -> dict[str, str]: - if not options: - return {} - aliases = { - "sep": "delimiter", - "delimiter": "delimiter", - "quote": "quote", - "escape": "escape", - "nullvalue": "nullValue", - "dateformat": "dateFormat", - "timestampformat": "timestampFormat", - "header": "header", - } - normalized: dict[str, str] = {} - for key, value in options.items(): - canonical = aliases.get(key.lower()) - if canonical is None: - continue - normalized[canonical] = str(value) - return normalized - - -def _option_lookup(options: Optional[Mapping[str, str]], *names: str) -> Optional[str]: - if not options: - return None - lowered = {key.lower(): value for key, value in options.items()} - for name in names: - if name.lower() in lowered: - return str(lowered[name.lower()]) - return None - - def _pandas_write_kwargs(options: Optional[Mapping[str, str]], header: bool) -> dict: kwargs: dict = {"index": False, "header": header} - delimiter = _option_lookup(options, "delimiter", "sep") + delimiter = option_lookup(options, "delimiter", "sep") if delimiter: kwargs["sep"] = delimiter - quote = _option_lookup(options, "quote") + quote = option_lookup(options, "quote") if quote: kwargs["quotechar"] = quote - escape = _option_lookup(options, "escape") + escape = option_lookup(options, "escape") if escape: kwargs["escapechar"] = escape kwargs["doublequote"] = False - null_value = _option_lookup(options, "nullValue") + null_value = option_lookup(options, "nullValue") if null_value is not None: kwargs["na_rep"] = null_value - date_format = _option_lookup(options, "dateFormat", "timestampFormat") + date_format = option_lookup(options, "dateFormat", "timestampFormat") if date_format: kwargs["date_format"] = java_date_format_to_strftime(date_format) return kwargs @@ -292,19 +276,19 @@ def _pandas_write_kwargs(options: Optional[Mapping[str, str]], header: bool) -> def _pandas_read_csv(csv_path: str, options: Optional[Mapping[str, str]]) -> pd.DataFrame: kwargs: dict = {} - delimiter = _option_lookup(options, "delimiter", "sep") + delimiter = option_lookup(options, "delimiter", "sep") if delimiter: kwargs["sep"] = delimiter - quote = _option_lookup(options, "quote") + quote = option_lookup(options, "quote") if quote: kwargs["quotechar"] = quote - escape = _option_lookup(options, "escape") + escape = option_lookup(options, "escape") if escape: kwargs["escapechar"] = escape - null_value = _option_lookup(options, "nullValue") + null_value = option_lookup(options, "nullValue") if null_value is not None: kwargs["na_values"] = [null_value] - header = _option_lookup(options, "header") + header = option_lookup(options, "header") if header and header.lower() == "false": kwargs["header"] = None return pd.read_csv(csv_path, **kwargs) @@ -345,26 +329,32 @@ def _format_value(value, strftime_fmt: str): def _format_temporal_columns(pdf: pd.DataFrame, options: Mapping[str, str]) -> pd.DataFrame: - date_fmt = _option_lookup(options, "dateFormat") - ts_fmt = _option_lookup(options, "timestampFormat") + date_fmt = option_lookup(options, "dateFormat") + ts_fmt = option_lookup(options, "timestampFormat") if not date_fmt and not ts_fmt: return pdf formatted = pdf.copy() for column in formatted.columns: series = formatted[column] if pd.api.types.is_datetime64_any_dtype(series): - has_time = bool((series.dt.hour.fillna(0) != 0).any() or (series.dt.minute.fillna(0) != 0).any()) - chosen = ts_fmt if (has_time and ts_fmt) else (date_fmt or ts_fmt) + has_time = bool( + (series.dt.hour.fillna(0) != 0).any() or (series.dt.minute.fillna(0) != 0).any() + ) + fallback = date_fmt or ts_fmt or "" + chosen = ts_fmt if (has_time and ts_fmt) else fallback formatted[column] = series.dt.strftime(java_date_format_to_strftime(chosen)) continue sample = series.dropna() if sample.empty or not hasattr(sample.iloc[0], "strftime"): continue - chosen = ts_fmt if ts_fmt and hasattr(sample.iloc[0], "hour") else (date_fmt or ts_fmt) + if ts_fmt and hasattr(sample.iloc[0], "hour"): + chosen = ts_fmt + else: + chosen = date_fmt or ts_fmt or "" if not chosen: continue strftime_fmt = java_date_format_to_strftime(chosen) - formatted[column] = series.map(lambda value: _format_value(value, strftime_fmt)) + formatted[column] = series.map(lambda value, fmt=strftime_fmt: _format_value(value, fmt)) return formatted @@ -378,14 +368,12 @@ def _spark_predicate_to_pandas_query(predicate: str) -> str: return query -def _apply_replace_where(existing_pdf: pd.DataFrame, predicate: str) -> pd.DataFrame: +def _drop_rows_matching(existing_pdf: pd.DataFrame, predicate: str) -> pd.DataFrame: query = _spark_predicate_to_pandas_query(predicate) try: matching = existing_pdf.query(query) except Exception as exc: - raise ValueError( - f"Unparsable replaceWhere predicate: {predicate!r}" - ) from exc + raise ValueError(f"Unparsable replaceWhere predicate: {predicate!r}") from exc return existing_pdf.drop(matching.index) diff --git a/src/testbricks/data_frame_wrapper.py b/src/testbricks/data_frame_wrapper.py index 076eeb9..d108f42 100644 --- a/src/testbricks/data_frame_wrapper.py +++ b/src/testbricks/data_frame_wrapper.py @@ -4,7 +4,7 @@ from pyspark.sql.group import GroupedData from pyspark.sql.utils import AnalysisException -from .catalog import TableIdentifier +from .catalog import TableIdentifier, normalize_csv_options, option_flag logger = logging.getLogger(__name__) @@ -29,6 +29,16 @@ def _resolve_file_format(fmt) -> str: return resolved +def _flatten_columns(cols): + flattened = [] + for col in cols: + if isinstance(col, (list, tuple)): + flattened.extend(col) + else: + flattened.append(col) + return flattened + + class _IoBuilder: """Shared format/option chaining used by both reader and writer.""" @@ -57,7 +67,7 @@ def __init__(self, spark_proxy): def table(self, table_name): ident = TableIdentifier.parse(table_name) merged = {"header": "true", "inferSchema": "true", **self._options} - df = self._spark._catalog.read_csv(ident, merged) + df = self._spark.read_table(ident, merged) return DataFrameWrapper(self._spark, df) @@ -71,24 +81,32 @@ def __init__(self, spark_proxy, dataframe): self._bucket_by = None self._sort_by = () + @classmethod + def _for_table_write( + cls, + spark_proxy, + dataframe, + *, + mode, + format=None, + options=None, + partition_by=(), + ): + writer = cls(spark_proxy, dataframe) + writer._mode = mode + writer._format = format + writer._partition_by = tuple(partition_by) + if options: + writer._options.update(options) + return writer + def partitionBy(self, *cols): - flattened = [] - for col in cols: - if isinstance(col, (list, tuple)): - flattened.extend(col) - else: - flattened.append(col) - self._partition_by = tuple(flattened) + self._partition_by = tuple(_flatten_columns(cols)) return self def bucketBy(self, numBuckets, *cols): """Accepted no-op: Hive-style bucketing is not simulated locally.""" - flattened = [] - for col in cols: - if isinstance(col, (list, tuple)): - flattened.extend(col) - else: - flattened.append(col) + flattened = _flatten_columns(cols) self._bucket_by = (numBuckets, tuple(flattened)) logger.info( "bucketBy(%s, %s) is accepted but not simulated", @@ -99,9 +117,8 @@ def bucketBy(self, numBuckets, *cols): def sortBy(self, col, *cols): """Accepted no-op: sortBy is not simulated locally.""" - flattened = [col, *cols] - self._sort_by = tuple(flattened) - logger.info("sortBy(%s) is accepted but not simulated", flattened) + self._sort_by = tuple([col, *cols]) + logger.info("sortBy(%s) is accepted but not simulated", [col, *cols]) return self def mode(self, save_mode): @@ -109,20 +126,20 @@ def mode(self, save_mode): return self def csv(self, path): - self._native_writer().csv(self._spark._get_full_path(path)) + self._native_writer().csv(self._spark.full_path(path)) def parquet(self, path): - self._native_writer().parquet(self._spark._get_full_path(path)) + self._native_writer().parquet(self._spark.full_path(path)) def json(self, path): - self._native_writer().json(self._spark._get_full_path(path)) + self._native_writer().json(self._spark.full_path(path)) def save(self, path, format=None, **options): if options: self._options.update(options) fmt = format or self._format or "parquet" resolved = _resolve_file_format(fmt) - full_path = self._spark._get_full_path(path) + full_path = self._spark.full_path(path) writer = self._native_writer() if resolved == "csv": writer.csv(full_path) @@ -158,49 +175,28 @@ def _header_flag(self): def _replace_where(self): return self._options.get("replaceWhere") or self._options.get("replacewhere") - def _csv_options(self): - return { - key: value - for key, value in self._options.items() - if key.lower() - in { - "delimiter", - "sep", - "quote", - "escape", - "nullvalue", - "dateformat", - "timestampformat", - "header", - } - } - - def _option_flag(self, *names): - wanted = {name.lower() for name in names} - for key, value in self._options.items(): - if key.lower() in wanted: - return str(value).lower() in {"true", "1", "yes"} - return False - - def saveAsTable(self, table_name): - ident = TableIdentifier.parse(table_name) - self._validate_partition_columns() - self._spark._catalog.save_dataframe( + def _save_to_table(self, ident, mode): + self._spark.save_table( ident, self._dataframe, - mode=self._mode, + mode=mode, header=self._header_flag(), replace_where=self._replace_where(), - csv_options=self._csv_options(), - overwrite_schema=self._option_flag("overwriteSchema"), - merge_schema=self._option_flag("mergeSchema"), + csv_options=normalize_csv_options(self._options), + overwrite_schema=option_flag(self._options, "overwriteSchema"), + merge_schema=option_flag(self._options, "mergeSchema"), ) + def saveAsTable(self, table_name): + ident = TableIdentifier.parse(table_name) + self._validate_partition_columns() + self._save_to_table(ident, self._mode) + def insertInto(self, table_name, overwrite=False): """Append or overwrite rows in an existing table (Spark DataFrameWriter.insertInto).""" ident = TableIdentifier.parse(table_name) self._validate_partition_columns() - if not self._spark._catalog.exists(ident): + if not self._spark.catalog.exists(ident): raise AnalysisException( f"[TABLE_OR_VIEW_NOT_FOUND] The table or view {ident} cannot be found. " "Verify the table exists before calling insertInto." @@ -210,19 +206,10 @@ def insertInto(self, table_name, overwrite=False): mode = "overwrite" else: mode = "append" - self._spark._catalog.save_dataframe( - ident, - self._dataframe, - mode=mode, - header=self._header_flag(), - replace_where=self._replace_where(), - csv_options=self._csv_options(), - overwrite_schema=self._option_flag("overwriteSchema"), - merge_schema=self._option_flag("mergeSchema"), - ) + self._save_to_table(ident, mode) -class DataFrameWriterV2: +class DataFrameWriterV2(_IoBuilder): """Minimal Spark DataFrameWriterV2 façade over ``saveAsTable``. Implements ``create`` / ``replace`` / ``append``. Full V2 verbs such as @@ -231,53 +218,46 @@ class DataFrameWriterV2: """ def __init__(self, spark_proxy, dataframe, table_name): + super().__init__() self._spark = spark_proxy self._dataframe = dataframe self._table_name = table_name - self._options = {} self._partitioned_by = () - self._using = None def using(self, provider): - self._using = provider - return self - - def option(self, key, value): - self._options[key] = value - return self - - def options(self, **kwargs): - self._options.update(kwargs) + self._format = provider return self def tableProperty(self, property, value): return self def partitionedBy(self, *cols): - flattened = [] - for col in cols: - if isinstance(col, (list, tuple)): - flattened.extend(col) - else: - flattened.append(col) - self._partitioned_by = tuple(flattened) + self._partitioned_by = tuple(_flatten_columns(cols)) return self def _writer(self, mode): - writer = DataFrameWriter(self._spark, self._dataframe) - writer._mode = mode - writer._format = self._using - writer._partition_by = self._partitioned_by - writer._options.update(self._options) - return writer + return DataFrameWriter._for_table_write( + self._spark, + self._dataframe, + mode=mode, + format=self._format, + options=self._options, + partition_by=self._partitioned_by, + ) def create(self): self._writer("error").saveAsTable(self._table_name) def replace(self): - writer = self._writer("overwrite") - writer._options.setdefault("overwriteSchema", "true") - writer.saveAsTable(self._table_name) + options = {"overwriteSchema": "true", **self._options} + DataFrameWriter._for_table_write( + self._spark, + self._dataframe, + mode="overwrite", + format=self._format, + options=options, + partition_by=self._partitioned_by, + ).saveAsTable(self._table_name) def append(self): self._writer("append").saveAsTable(self._table_name) diff --git a/src/testbricks/dbutils/__init__.py b/src/testbricks/dbutils/__init__.py index f2e8895..1a775a3 100644 --- a/src/testbricks/dbutils/__init__.py +++ b/src/testbricks/dbutils/__init__.py @@ -1,5 +1,5 @@ from .dbutils_mock import DbutilsMock -from .errors import DbutilsError +from .errors import DbutilsError as DbutilsError dbutils = DbutilsMock() diff --git a/src/testbricks/dbutils/dbutils_mock.py b/src/testbricks/dbutils/dbutils_mock.py index 548b385..4b00f25 100644 --- a/src/testbricks/dbutils/dbutils_mock.py +++ b/src/testbricks/dbutils/dbutils_mock.py @@ -11,15 +11,13 @@ class DbutilsMock: def __init__(self): - from testbricks.notebook_executor import NotebookExecutor - self._path_resolver = PathResolver() self._source_dir = None - self._executor = NotebookExecutor(self) + self._notebook_executor = None self.fs = FsMock(self._path_resolver) self.secrets = SecretsMock() self.widgets = WidgetsMock() - self.notebook = NotebookMock(self._executor) + self.notebook = NotebookMock(self) self.jobs = JobsMock() self.library = LibraryMock() self.data = DataMock() @@ -28,16 +26,25 @@ def __init__(self): def source_dir(self): return self._source_dir + @property + def path_resolver(self): + return self._path_resolver + @property def executor(self): - return self._executor + # Lazy seam: notebook_executor imports dbutils submodules at module + # scope, so constructing the executor here would create an import + # cycle (dbutils package -> dbutils_mock -> notebook_executor -> + # dbutils package). Importing on first use breaks the cycle. + if self._notebook_executor is None: + from testbricks.notebook_executor import NotebookExecutor + + self._notebook_executor = NotebookExecutor(self) + return self._notebook_executor def configure(self, base_path, source_dir=None): - self._path_resolver.configure(base_path) + self._path_resolver.configure(base_path, source_dir) self._source_dir = source_dir - def help(self, module=None): - return True - def __getattr__(self, name): return NoOpModule() diff --git a/src/testbricks/dbutils/errors.py b/src/testbricks/dbutils/errors.py index f1e1cf4..86e185e 100644 --- a/src/testbricks/dbutils/errors.py +++ b/src/testbricks/dbutils/errors.py @@ -1,2 +1,5 @@ -class DbutilsError(Exception): +from testbricks.errors import TestbricksError + + +class DbutilsError(TestbricksError): """Raised when a dbutils mock operation fails.""" diff --git a/src/testbricks/dbutils/fs.py b/src/testbricks/dbutils/fs.py index 1f07444..cf18ceb 100644 --- a/src/testbricks/dbutils/fs.py +++ b/src/testbricks/dbutils/fs.py @@ -16,30 +16,24 @@ def __init__(self, path_resolver: PathResolver): def __getattr__(self, name): return NoOpModule() - def help(self, command=None): - return True - def cp(self, from_path, to_path, recurse=False): source = self._resolve(from_path) destination = self._resolve(to_path) + source_is_dir = os.path.isdir(source) def _copy(): - if os.path.isdir(source): + if source_is_dir: if not recurse: - raise DbutilsError( - f"source is a directory and recurse is False: {from_path}" - ) + raise DbutilsError(f"source is a directory and recurse is False: {from_path}") if os.path.exists(destination): raise DbutilsError(f"destination already exists: {to_path}") self._ensure_parent(destination) - if os.path.isdir(source): + if source_is_dir: shutil.copytree(source, destination) else: shutil.copy2(source, destination) - return self._os_call( - f"failed to copy {from_path} to {to_path}", _copy - ) + return self._os_call(f"failed to copy {from_path} to {to_path}", _copy) def mv(self, from_path, to_path, recurse=False): self.cp(from_path, to_path, recurse=recurse) @@ -55,7 +49,7 @@ def _remove(): elif os.path.isdir(target): shutil.rmtree(target) if recurse else os.rmdir(target) else: - raise FileNotFoundError(target) + raise DbutilsError(f"failed to remove {path}") return self._os_call(f"failed to remove {path}", _remove) @@ -122,6 +116,6 @@ def _os_call(message, action): action() except DbutilsError: raise - except (OSError, FileNotFoundError) as exc: + except OSError as exc: raise DbutilsError(message) from exc return True diff --git a/src/testbricks/dbutils/jobs.py b/src/testbricks/dbutils/jobs.py index 11a242c..b0e6dd9 100644 --- a/src/testbricks/dbutils/jobs.py +++ b/src/testbricks/dbutils/jobs.py @@ -5,9 +5,7 @@ from .errors import DbutilsError _MISSING = object() -_current_task_key: ContextVar[str | None] = ContextVar( - "task_values_current_task", default=None -) +_current_task_key: ContextVar[str | None] = ContextVar("task_values_current_task", default=None) def _stringify(value): @@ -97,9 +95,7 @@ def get( return _stringify(debugValue) if default is not _MISSING: return _stringify(default) - raise DbutilsError( - f"Task value '{key}' not found for task '{resolved_task}'" - ) + raise DbutilsError(f"Task value '{key}' not found for task '{resolved_task}'") class JobsMock: diff --git a/src/testbricks/dbutils/notebook.py b/src/testbricks/dbutils/notebook.py index 533c4c1..40152e4 100644 --- a/src/testbricks/dbutils/notebook.py +++ b/src/testbricks/dbutils/notebook.py @@ -2,11 +2,15 @@ class NotebookMock: - def __init__(self, executor): - self._executor = executor + def __init__(self, dbutils_mock): + self._dbutils = dbutils_mock + + @property + def _executor(self): + return self._dbutils.executor def exit(self, value): raise NotebookExit(str(value)) def run(self, path, timeout_seconds, arguments=None): - return self._executor.run_isolated(path, arguments=arguments or {}) + return self._executor.run_isolated(path, arguments=arguments) diff --git a/src/testbricks/dbutils/path_resolver.py b/src/testbricks/dbutils/path_resolver.py index 571036c..31787fc 100644 --- a/src/testbricks/dbutils/path_resolver.py +++ b/src/testbricks/dbutils/path_resolver.py @@ -2,7 +2,8 @@ from .errors import DbutilsError -_PREFIXES = ("dbfs:/", "/dbfs/", "/mnt/") +DBFS_PREFIXES = ("dbfs:/", "/dbfs/", "/mnt/") +WORKSPACE_PREFIXES = ("/Workspace/", "/Repos/") def strip_known_prefix(path, prefixes): @@ -13,23 +14,27 @@ def strip_known_prefix(path, prefixes): class PathResolver: + """Resolves ``dbutils.fs`` paths (against ``base_path``) and notebook paths + (against ``source_dir`` or the calling notebook) into local filesystem paths.""" + def __init__(self): self._base_path = None + self._source_dir = None - def configure(self, base_path): - self._base_path = os.path.abspath(base_path) + def configure(self, base_path, source_dir=None): + self._base_path = os.path.abspath(base_path) if base_path is not None else None + self._source_dir = os.path.abspath(source_dir) if source_dir else None @property def is_configured(self): return self._base_path is not None def resolve(self, path): + """Resolve a dbutils.fs-style path under the configured base_path.""" if not self.is_configured: - raise DbutilsError( - "dbutils not configured — call configure(base_path) first" - ) + raise DbutilsError("dbutils not configured — call configure(base_path) first") - remainder = strip_known_prefix(path, _PREFIXES).lstrip("/") + remainder = strip_known_prefix(path, DBFS_PREFIXES).lstrip("/") resolved = os.path.normpath(os.path.join(self._base_path, remainder)) base = os.path.normpath(self._base_path) @@ -37,3 +42,35 @@ def resolve(self, path): raise DbutilsError(f"path escapes base_path: {path}") return resolved + + def resolve_notebook(self, path, caller_file=None): + """Resolve a %run / notebook.run-style path. + + Workspace-style absolute paths (``/Workspace/...``, ``/Repos/...``) are + rooted at the configured ``source_dir`` (with a containment check); + relative paths are resolved against the calling notebook's directory. + """ + if path.startswith("/"): + if self._source_dir is None: + raise DbutilsError("source_dir not configured — required for workspace paths") + remainder = strip_known_prefix(path, WORKSPACE_PREFIXES).lstrip("/") + notebook_path = self._require_within_source_dir( + os.path.join(self._source_dir, remainder), path + ) + else: + if not caller_file: + raise DbutilsError("caller file not set — cannot resolve relative notebook path") + notebook_path = os.path.normpath(os.path.join(os.path.dirname(caller_file), path)) + + if not notebook_path.endswith(".py"): + notebook_path += ".py" + if not os.path.exists(notebook_path): + raise DbutilsError(f"Notebook not found: {notebook_path}") + return notebook_path + + def _require_within_source_dir(self, notebook_path, original_path): + resolved = os.path.normpath(notebook_path) + source_dir = os.path.normpath(self._source_dir) + if resolved != source_dir and not resolved.startswith(source_dir + os.sep): + raise DbutilsError(f"path escapes source_dir: {original_path}") + return resolved diff --git a/src/testbricks/dbutils/secrets.py b/src/testbricks/dbutils/secrets.py index 0c0b2fd..cb2805e 100644 --- a/src/testbricks/dbutils/secrets.py +++ b/src/testbricks/dbutils/secrets.py @@ -29,16 +29,16 @@ class SecretsMock: def get(self, scope, key): env_name = _secret_env_name(scope, key) if env_name not in os.environ: - raise DbutilsError( - f"Secret for scope '{scope}' and key '{key}' does not exist" - ) + raise DbutilsError(f"Secret for scope '{scope}' and key '{key}' does not exist") return os.environ[env_name] def getBytes(self, scope, key): return self.get(scope, key).encode("utf-8") def list(self, scope): - keys = sorted({key for listed_scope, key in _iter_secret_entries() if listed_scope == scope}) + keys = sorted( + {key for listed_scope, key in _iter_secret_entries() if listed_scope == scope} + ) return [SecretMetadata(key=key) for key in keys] def listScopes(self): diff --git a/src/testbricks/dbutils/widgets.py b/src/testbricks/dbutils/widgets.py index a5843d8..66ee347 100644 --- a/src/testbricks/dbutils/widgets.py +++ b/src/testbricks/dbutils/widgets.py @@ -18,6 +18,27 @@ def argument_override_context(keys): _argument_overrides.reset(token) +@contextmanager +def seeded_environ(values, *, overwrite=True): + """Set env vars for the duration of the block, then restore prior state. + + With ``overwrite=False`` existing env vars win (used for workflow + ``base_parameters``, which must not clobber real environment settings). + """ + saved = {key: os.environ.get(key) for key in values} + for key, value in values.items(): + if overwrite or key not in os.environ: + os.environ[key] = str(value) + try: + yield + finally: + for key, original in saved.items(): + if original is None: + os.environ.pop(key, None) + else: + os.environ[key] = original + + class WidgetsMock: def __init__(self): self._registry: set[str] = set() @@ -37,9 +58,7 @@ def text(self, name, default, label=None): def dropdown(self, name, default, choices, label=None): if default not in choices: - raise DbutilsError( - f"Default value '{default}' is not in choices {list(choices)}" - ) + raise DbutilsError(f"Default value '{default}' is not in choices {list(choices)}") return self._register(name, default) def combobox(self, name, default, choices, label=None): @@ -52,9 +71,7 @@ def multiselect(self, name, default, choices, label=None): selected = [default] for value in selected: if value not in choices: - raise DbutilsError( - f"Default value '{default}' is not in choices {list(choices)}" - ) + raise DbutilsError(f"Default value '{default}' is not in choices {list(choices)}") return self._register(name, default) def get(self, name): @@ -66,9 +83,7 @@ def getAll(self): return {name: os.environ[name] for name in self._registry} def getArgument(self, name, optional=None): - if name in self._registry: - return self.get(name) - if optional is not None: + if name not in self._registry and optional is not None: return str(optional) return self.get(name) diff --git a/src/testbricks/errors.py b/src/testbricks/errors.py new file mode 100644 index 0000000..de3e20e --- /dev/null +++ b/src/testbricks/errors.py @@ -0,0 +1,5 @@ +"""Common base for all testbricks exception hierarchies.""" + + +class TestbricksError(Exception): + """Base class for every exception raised by the testbricks mocks.""" diff --git a/src/testbricks/local_workflow_runner.py b/src/testbricks/local_workflow_runner.py index 3ebe769..57e48ae 100644 --- a/src/testbricks/local_workflow_runner.py +++ b/src/testbricks/local_workflow_runner.py @@ -2,9 +2,11 @@ import os import re import time -from contextlib import contextmanager +from dataclasses import dataclass, field from graphlib import CycleError, TopologicalSorter +from testbricks.dbutils import configure, dbutils +from testbricks.dbutils.widgets import argument_override_context, seeded_environ from testbricks.notebook_executor import transform_run_commands __all__ = ["LocalWorkflowRunner", "transform_run_commands"] @@ -38,6 +40,196 @@ def _require(condition, message): raise ValueError(message) +@dataclass(frozen=True) +class TaskSpec: + """Validated, kind-agnostic view of one workflow task.""" + + task_key: str + kind: str # "notebook" | "condition" | "for_each" + notebook_name: str | None = None + base_parameters: dict = field(default_factory=dict) + depends_on: list = field(default_factory=list) + dep_specs: list = field(default_factory=list) + run_if: str = "ALL_SUCCESS" + condition_task: dict | None = None + for_each_task: dict | None = None + max_retries: int = 0 + retry_interval_ms: int = 0 + timeout_seconds: int = 0 + + +def _extract_notebook_name(notebook_path, task_key): + _require( + isinstance(notebook_path, str) and notebook_path.strip(), + f"Task '{task_key}' has an invalid notebook_path", + ) + notebook_name = notebook_path.rstrip("/").split("/")[-1] + _require(notebook_name, f"Task '{task_key}' has an invalid notebook_path") + return notebook_name + + +def _parse_non_negative_int(raw, default, field_name, task_key): + if raw is None: + return default + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"Task '{task_key}' has invalid '{field_name}': {raw!r}") from exc + _require(value >= 0, f"Task '{task_key}' has invalid '{field_name}': {raw!r}") + return value + + +def _parse_condition_task(condition_task, task_key): + _require( + condition_task.get("op"), + f"Task '{task_key}' has invalid 'condition_task'", + ) + op = str(condition_task.get("op")).upper() + _require( + op in CONDITION_OPS, + f"Unsupported condition op '{condition_task.get('op')}' on task '{task_key}'", + ) + _require( + "left" in condition_task and "right" in condition_task, + f"Task '{task_key}' condition_task requires 'left' and 'right'", + ) + return {**condition_task, "op": op} + + +def _parse_for_each_task(for_each_task, task_key): + nested = for_each_task.get("task") + _require( + isinstance(nested, dict), + f"Task '{task_key}' for_each_task requires a nested 'task'", + ) + nested_notebook = nested.get("notebook_task") + _require( + isinstance(nested_notebook, dict), + f"Task '{task_key}' for_each nested task is missing 'notebook_task'", + ) + _extract_notebook_name(nested_notebook.get("notebook_path"), task_key) + nested_params = nested_notebook.get("base_parameters", {}) + _require( + isinstance(nested_params, dict), + f"Task '{task_key}' has invalid nested 'base_parameters' format", + ) + _require( + "inputs" in for_each_task, + f"Task '{task_key}' for_each_task requires 'inputs'", + ) + return for_each_task + + +def _parse_dependencies(depends_on, task_key): + _require( + isinstance(depends_on, list), + f"Task '{task_key}' has invalid 'depends_on' format", + ) + dependency_keys = [] + dep_specs = [] + for dependency in depends_on: + _require( + isinstance(dependency, dict) and dependency.get("task_key"), + f"Task '{task_key}' has malformed dependency entry", + ) + dependency_keys.append(dependency["task_key"]) + dep_specs.append((dependency["task_key"], dependency.get("outcome"))) + return dependency_keys, dep_specs + + +def _parse_run_if(run_if, task_key): + run_if = run_if or "ALL_SUCCESS" + _require( + isinstance(run_if, str) and run_if.upper() in RUN_IF_VALUES, + f"Unsupported run_if '{run_if}' on task '{task_key}'", + ) + return run_if.upper() + + +def parse_workflow(tasks) -> list[TaskSpec]: + """Validate a workflow ``tasks`` list and return one TaskSpec per task.""" + specs = [] + seen_task_keys = set() + seen_notebook_names = set() + for task in tasks: + _require(isinstance(task, dict), "Each task must be a JSON object") + task_key = task.get("task_key") + _require(task_key, "Each task must include a non-empty 'task_key'") + _require(task_key not in seen_task_keys, f"Duplicate task_key found: {task_key}") + seen_task_keys.add(task_key) + + notebook_task = task.get("notebook_task") + condition_task = task.get("condition_task") + for_each_task = task.get("for_each_task") + has_notebook = isinstance(notebook_task, dict) + has_condition = isinstance(condition_task, dict) + has_for_each = isinstance(for_each_task, dict) + _require( + has_notebook or has_condition or has_for_each, + f"Task '{task_key}' is missing 'notebook_task'", + ) + _require( + [has_notebook, has_condition, has_for_each].count(True) == 1, + f"Task '{task_key}' cannot mix notebook, condition, and for_each tasks", + ) + + notebook_name = None + base_parameters = {} + if has_notebook: + notebook_name = _extract_notebook_name(notebook_task.get("notebook_path"), task_key) + _require( + notebook_name not in seen_notebook_names, + f"Duplicate notebook name found: {notebook_name}", + ) + seen_notebook_names.add(notebook_name) + base_parameters = notebook_task.get("base_parameters", {}) + _require( + isinstance(base_parameters, dict), + f"Task '{task_key}' has invalid 'base_parameters' format", + ) + elif has_condition: + condition_task = _parse_condition_task(condition_task, task_key) + else: + for_each_task = _parse_for_each_task(for_each_task, task_key) + + depends_on, dep_specs = _parse_dependencies(task.get("depends_on", []), task_key) + run_if = _parse_run_if(task.get("run_if"), task_key) + + if has_condition: + kind = "condition" + elif has_for_each: + kind = "for_each" + else: + kind = "notebook" + + specs.append( + TaskSpec( + task_key=task_key, + kind=kind, + notebook_name=notebook_name, + base_parameters=base_parameters, + depends_on=depends_on, + dep_specs=dep_specs, + run_if=run_if, + condition_task=condition_task, + for_each_task=for_each_task, + max_retries=_parse_non_negative_int( + task.get("max_retries"), 0, "max_retries", task_key + ), + retry_interval_ms=_parse_non_negative_int( + task.get("min_retry_interval_millis"), + 0, + "min_retry_interval_millis", + task_key, + ), + timeout_seconds=_parse_non_negative_int( + task.get("timeout_seconds"), 0, "timeout_seconds", task_key + ), + ) + ) + return specs + + def _matches_run_if(run_if, statuses): if not statuses: return True @@ -127,157 +319,24 @@ def __init__(self, source_dir, workflow_json_path, base_path): raise ValueError("Workflow graph contains a cycle") from exc def _parse_tasks(self, tasks): - for task in tasks: - _require(isinstance(task, dict), "Each task must be a JSON object") - task_key = task.get("task_key") - _require(task_key, "Each task must include a non-empty 'task_key'") - _require( - task_key not in self._task_kind, - f"Duplicate task_key found: {task_key}", - ) - - notebook_task = task.get("notebook_task") - condition_task = task.get("condition_task") - for_each_task = task.get("for_each_task") - has_notebook = isinstance(notebook_task, dict) - has_condition = isinstance(condition_task, dict) - has_for_each = isinstance(for_each_task, dict) - _require( - has_notebook or has_condition or has_for_each, - f"Task '{task_key}' is missing 'notebook_task'", - ) - _require( - [has_notebook, has_condition, has_for_each].count(True) == 1, - f"Task '{task_key}' cannot mix notebook, condition, and for_each tasks", - ) - - notebook_name = None - base_parameters = {} - if has_notebook: - notebook_name = self._extract_notebook_name( - notebook_task.get("notebook_path"), task_key - ) - _require( - notebook_name not in self._notebook_insertion_order, - f"Duplicate notebook name found: {notebook_name}", - ) - base_parameters = notebook_task.get("base_parameters", {}) - _require( - isinstance(base_parameters, dict), - f"Task '{task_key}' has invalid 'base_parameters' format", - ) - elif has_condition: - _require( - condition_task.get("op"), - f"Task '{task_key}' has invalid 'condition_task'", - ) - op = str(condition_task.get("op")).upper() - _require( - op in CONDITION_OPS, - f"Unsupported condition op '{condition_task.get('op')}' on task '{task_key}'", - ) - _require( - "left" in condition_task and "right" in condition_task, - f"Task '{task_key}' condition_task requires 'left' and 'right'", - ) - condition_task = {**condition_task, "op": op} - else: - nested = for_each_task.get("task") - _require( - isinstance(nested, dict), - f"Task '{task_key}' for_each_task requires a nested 'task'", - ) - nested_notebook = nested.get("notebook_task") - _require( - isinstance(nested_notebook, dict), - f"Task '{task_key}' for_each nested task is missing 'notebook_task'", - ) - self._extract_notebook_name( - nested_notebook.get("notebook_path"), task_key - ) - nested_params = nested_notebook.get("base_parameters", {}) - _require( - isinstance(nested_params, dict), - f"Task '{task_key}' has invalid nested 'base_parameters' format", - ) - _require( - "inputs" in for_each_task, - f"Task '{task_key}' for_each_task requires 'inputs'", - ) - - depends_on = task.get("depends_on", []) - _require( - isinstance(depends_on, list), - f"Task '{task_key}' has invalid 'depends_on' format", - ) - dependency_keys = [] - dep_specs = [] - for dependency in depends_on: - _require( - isinstance(dependency, dict) and dependency.get("task_key"), - f"Task '{task_key}' has malformed dependency entry", - ) - dependency_keys.append(dependency["task_key"]) - dep_specs.append((dependency["task_key"], dependency.get("outcome"))) - - run_if = task.get("run_if") or "ALL_SUCCESS" - _require( - isinstance(run_if, str) and run_if.upper() in RUN_IF_VALUES, - f"Unsupported run_if '{run_if}' on task '{task_key}'", - ) - - if has_condition: - kind = "condition" - elif has_for_each: - kind = "for_each" - else: - kind = "notebook" - self._task_kind[task_key] = kind - self._condition_task[task_key] = condition_task if has_condition else None - self._for_each_task[task_key] = for_each_task if has_for_each else None - self._task_to_notebook[task_key] = notebook_name - if notebook_name is not None: - self._notebook_to_task[notebook_name] = task_key - self._notebook_insertion_order.append(notebook_name) - self._notebook_base_params[notebook_name] = base_parameters - self._task_dependencies[task_key] = dependency_keys - self._task_dep_specs[task_key] = dep_specs - self._task_run_if[task_key] = run_if.upper() - self._task_max_retries[task_key] = self._parse_non_negative_int( - task.get("max_retries"), 0, "max_retries", task_key - ) - self._task_retry_interval_ms[task_key] = self._parse_non_negative_int( - task.get("min_retry_interval_millis"), - 0, - "min_retry_interval_millis", - task_key, - ) - self._task_timeout_seconds[task_key] = self._parse_non_negative_int( - task.get("timeout_seconds"), 0, "timeout_seconds", task_key - ) + for spec in parse_workflow(tasks): + task_key = spec.task_key + self._task_kind[task_key] = spec.kind + self._condition_task[task_key] = spec.condition_task + self._for_each_task[task_key] = spec.for_each_task + self._task_to_notebook[task_key] = spec.notebook_name + if spec.notebook_name is not None: + self._notebook_to_task[spec.notebook_name] = task_key + self._notebook_insertion_order.append(spec.notebook_name) + self._notebook_base_params[spec.notebook_name] = spec.base_parameters + self._task_dependencies[task_key] = spec.depends_on + self._task_dep_specs[task_key] = spec.dep_specs + self._task_run_if[task_key] = spec.run_if + self._task_max_retries[task_key] = spec.max_retries + self._task_retry_interval_ms[task_key] = spec.retry_interval_ms + self._task_timeout_seconds[task_key] = spec.timeout_seconds self._task_insertion_order.append(task_key) - self._task_base_params[task_key] = base_parameters - - def _extract_notebook_name(self, notebook_path, task_key): - _require( - isinstance(notebook_path, str) and notebook_path.strip(), - f"Task '{task_key}' has an invalid notebook_path", - ) - notebook_name = notebook_path.rstrip("/").split("/")[-1] - _require(notebook_name, f"Task '{task_key}' has an invalid notebook_path") - return notebook_name - - def _parse_non_negative_int(self, raw, default, field, task_key): - if raw is None: - return default - try: - value = int(raw) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Task '{task_key}' has invalid '{field}': {raw!r}" - ) from exc - _require(value >= 0, f"Task '{task_key}' has invalid '{field}': {raw!r}") - return value + self._task_base_params[task_key] = spec.base_parameters def _build_graphs(self): successors = {key: set() for key in self._task_insertion_order} @@ -309,9 +368,7 @@ def _is_eligible(self, task_key): status = self.task_statuses.get(dep_key) if status is None: return False - if outcome is not None and not self._outcome_matches( - dep_key, status, outcome - ): + if outcome is not None and not self._outcome_matches(dep_key, status, outcome): return False statuses = [self.task_statuses[dep_key] for dep_key, _ in specs] return _matches_run_if(self._task_run_if[task_key], statuses) @@ -373,14 +430,10 @@ def _selected_tasks(self, only, from_task): stack.append(successor) return selected - def _run_task_with_retries( - self, task_key, executor, store, execution_globals, extra_globals - ): + def _run_task_with_retries(self, task_key, executor, store, execution_globals, extra_globals): timeout = self._task_timeout_seconds.get(task_key) or 0 if timeout: - print( - f"Task '{task_key}' timeout_seconds={timeout} accepted but not enforced" - ) + print(f"Task '{task_key}' timeout_seconds={timeout} accepted but not enforced") def action(): kind = self._task_kind[task_key] @@ -403,15 +456,12 @@ def action(): if attempt > retries: raise print( - f"Task '{task_key}' failed " - f"(attempt {attempt}/{retries + 1}), retrying: {exc}" + f"Task '{task_key}' failed (attempt {attempt}/{retries + 1}), retrying: {exc}" ) if interval_ms: time.sleep(interval_ms / 1000.0) def _executor(self): - from testbricks.dbutils import dbutils - return dbutils.executor def _run_notebook(self, relative_path, namespace): @@ -420,30 +470,15 @@ def _run_notebook(self, relative_path, namespace): def _execfile(self, file_path, global_namespace, local_namespace): self._executor().exec_file(file_path, global_namespace, top_level=False) - @contextmanager - def _seeded_env(self, values): - saved = {key: os.environ.get(key) for key in values} - for key, value in values.items(): - os.environ.setdefault(key, str(value)) - try: - yield - finally: - for key, original in saved.items(): - if original is None: - os.environ.pop(key, None) - else: - os.environ[key] = original - def _run_notebook_task(self, task_key, executor, store, execution_globals): notebook_name = self._task_to_notebook[task_key] notebook_path = os.path.join(self.source_dir, f"{notebook_name}.py") if not os.path.exists(notebook_path): raise FileNotFoundError(f"Notebook file not found: {notebook_path}") - from testbricks.dbutils.widgets import argument_override_context base_params = self._task_base_params.get(task_key, {}) with ( - self._seeded_env(base_params), + seeded_environ(base_params, overwrite=False), argument_override_context(base_params.keys()), store.current_task(task_key), ): @@ -494,9 +529,7 @@ def _run_for_each_task(self, task_key, executor, store, extra_globals): inputs = self._resolve_for_each_inputs(spec.get("inputs"), store) nested = spec["task"] notebook_task = nested["notebook_task"] - notebook_name = self._extract_notebook_name( - notebook_task.get("notebook_path"), task_key - ) + notebook_name = _extract_notebook_name(notebook_task.get("notebook_path"), task_key) nested_key = nested.get("task_key") or task_key base_parameters = notebook_task.get("base_parameters", {}) or {} caller = os.path.join(self.source_dir, "_workflow.py") @@ -505,27 +538,17 @@ def _run_for_each_task(self, task_key, executor, store, extra_globals): key: self._render_input_template(value, item, index) for key, value in base_parameters.items() } - saved_env = {key: os.environ.get(key) for key in params} - try: - with ( - store.current_task(nested_key), - executor.caller_context(caller), - ): - executor.run_isolated( - f"/Workspace/{notebook_name}", - arguments=params, - extra=extra_globals, - ) - finally: - for key, original in saved_env.items(): - if original is None: - os.environ.pop(key, None) - else: - os.environ[key] = original + with ( + store.current_task(nested_key), + executor.caller_context(caller), + ): + executor.run_isolated( + f"/Workspace/{notebook_name}", + arguments=params, + extra=extra_globals, + ) def run_workflow(self, extra_globals=None, only=None, from_task=None): - from testbricks.dbutils import configure, dbutils - configure(self.base_path, source_dir=self.source_dir) executor = dbutils.executor store = dbutils.jobs.taskValues diff --git a/src/testbricks/notebook_exceptions.py b/src/testbricks/notebook_exceptions.py index 3a6cf1e..2acbeb9 100644 --- a/src/testbricks/notebook_exceptions.py +++ b/src/testbricks/notebook_exceptions.py @@ -1,10 +1,19 @@ +from testbricks.errors import TestbricksError + + class NotebookExit(BaseException): + """Notebook-level exit raised by ``dbutils.notebook.exit``. + + Deliberately derives from ``BaseException`` so that user notebook code + (and retry loops) cannot swallow it with a plain ``except Exception``. + """ + def __init__(self, value: str): self.value = value super().__init__(value) -class ShellCommandError(RuntimeError): +class ShellCommandError(TestbricksError, RuntimeError): def __init__(self, message, returncode=None): self.returncode = returncode super().__init__(message) diff --git a/src/testbricks/notebook_executor.py b/src/testbricks/notebook_executor.py index aebea49..6e2aba5 100644 --- a/src/testbricks/notebook_executor.py +++ b/src/testbricks/notebook_executor.py @@ -6,7 +6,9 @@ from contextlib import contextmanager from contextvars import ContextVar -from testbricks.dbutils.path_resolver import strip_known_prefix +from testbricks.catalog.identifier import strip_wrappers +from testbricks.dbutils.errors import DbutilsError +from testbricks.dbutils.widgets import argument_override_context, seeded_environ from testbricks.notebook_exceptions import NotebookExit, ShellCommandError RUN_COMMAND_PATTERN = re.compile( @@ -16,22 +18,18 @@ SH_START_PATTERN = re.compile(r"^(\s*)#\s*(?:MAGIC\s+)?%sh(?:\s+(.*))?\s*$") FS_START_PATTERN = re.compile(r"^(\s*)#\s*(?:MAGIC\s+)?%fs(?:\s+(.*))?\s*$") MAGIC_BODY_PATTERN = re.compile(r"^\s*#\s*MAGIC\s+(.*)$") -WORKSPACE_PREFIXES = ("/Workspace/", "/Repos/") _caller_file: ContextVar[str | None] = ContextVar("caller_file", default=None) def _strip_matching_quotes(value): - value = value.strip() - if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'): - return value[1:-1].strip() - return value + return strip_wrappers(value, "'\"").strip() def parse_run_path(raw_path, file_path): path = _strip_matching_quotes(raw_path) if not path: - raise ValueError(f"Empty %run path in notebook '{file_path}'") + raise DbutilsError(f"Empty %run path in notebook '{file_path}'") return path @@ -61,6 +59,29 @@ def _parse_sh_remainder(remainder): return fail_on_error, inline.strip() +def _collect_magic_lines(lines, index): + """Collect consecutive ``# MAGIC `` lines starting at ``index``. + + Stops at the first non-magic line or at a body that starts another magic + command (``%run``, ``%sh``, ...). Returns ``(bodies, next_index)``. + """ + bodies = [] + while index < len(lines): + body_match = MAGIC_BODY_PATTERN.match(lines[index].rstrip("\n")) + if not body_match: + break + body = body_match.group(1) + if body.lstrip().startswith("%"): + break + bodies.append(body) + index += 1 + return bodies, index + + +def _newline_for(line): + return "\n" if line.endswith("\n") else "" + + def transform_sh_commands(source): lines = source.splitlines(keepends=True) output = [] @@ -78,20 +99,13 @@ def transform_sh_commands(source): started_with_magic = "MAGIC" in line.split("%sh", 1)[0] index += 1 if started_with_magic: - while index < len(lines): - body_match = MAGIC_BODY_PATTERN.match(lines[index].rstrip("\n")) - if not body_match: - break - body = body_match.group(1) - if body.lstrip().startswith("%"): - break - script_lines.append(body) - index += 1 + bodies, index = _collect_magic_lines(lines, index) + script_lines.extend(bodies) script = "\n".join(script_lines) if script.strip(): - newline = "\n" if line.endswith("\n") else "" output.append( - f"{indent}__run_shell__({script!r}, fail_on_error={fail_on_error}){newline}" + f"{indent}__run_shell__({script!r}, fail_on_error={fail_on_error})" + f"{_newline_for(line)}" ) return "".join(output) @@ -125,26 +139,14 @@ def transform_fs_commands(source): parts = shlex.split(remainder or "") started_with_magic = "MAGIC" in line.split("%fs", 1)[0] index += 1 - extra = [] if started_with_magic: - while index < len(lines): - body_match = MAGIC_BODY_PATTERN.match(lines[index].rstrip("\n")) - if not body_match: - break - body = body_match.group(1) - if body.lstrip().startswith("%"): - break - extra.append(body) - index += 1 - if extra: - extra_text = "\n".join(extra) - if extra_text.strip(): - parts.extend(shlex.split(extra_text)) + bodies, index = _collect_magic_lines(lines, index) + if bodies: + parts.extend(shlex.split("\n".join(bodies))) if not parts: continue command, *args = parts - newline = "\n" if line.endswith("\n") else "" - output.append(f"{indent}{_fs_python_call(command, args)}{newline}") + output.append(f"{indent}{_fs_python_call(command, args)}{_newline_for(line)}") return "".join(output) @@ -197,31 +199,10 @@ def namespace(self, file_path, extra=None): return ns def resolve_path(self, path, caller_file=None): - from testbricks.dbutils.errors import DbutilsError - - normalized_path = _strip_matching_quotes(path) - if normalized_path.startswith("/"): - source_dir = self._dbutils.source_dir - if source_dir is None: - raise DbutilsError( - "source_dir not configured — required for workspace paths" - ) - remainder = strip_known_prefix(normalized_path, WORKSPACE_PREFIXES).lstrip("/") - notebook_path = os.path.join(source_dir, remainder) - else: - caller_file = caller_file or _caller_file.get() - if not caller_file: - raise DbutilsError( - "caller file not set — cannot resolve relative notebook path" - ) - notebook_path = os.path.join(os.path.dirname(caller_file), normalized_path) - - notebook_path = os.path.normpath(notebook_path) - if not notebook_path.endswith(".py"): - notebook_path += ".py" - if not os.path.exists(notebook_path): - raise DbutilsError(f"Notebook not found: {notebook_path}") - return notebook_path + return self._dbutils.path_resolver.resolve_notebook( + _strip_matching_quotes(path), + caller_file if caller_file is not None else _caller_file.get(), + ) def exec_file(self, file_path, namespace, *, top_level=False): with self.caller_context(file_path): @@ -245,16 +226,13 @@ def run_shared(self, path, namespace): self.exec_file(notebook_path, namespace, top_level=False) def run_isolated(self, path, arguments=None, extra=None): - from testbricks.dbutils.widgets import argument_override_context - arguments = arguments or {} notebook_path = self.resolve_path(path, caller_file=_caller_file.get()) - for key, value in arguments.items(): - os.environ[key] = str(value) namespace = self.namespace(notebook_path, extra=extra) with ( argument_override_context(arguments.keys()), + seeded_environ(arguments), self._dbutils.jobs.taskValues.isolated_context(), ): try: diff --git a/src/testbricks/spark_proxy.py b/src/testbricks/spark_proxy.py index 0ea9090..6012ff9 100644 --- a/src/testbricks/spark_proxy.py +++ b/src/testbricks/spark_proxy.py @@ -46,7 +46,35 @@ def sql(self, query): return self._wrap(self._spark_session.sql(rewrite_from_join_identifiers(query))) def parallelize(self, c, numSlices=None): - return self._wrap(self._spark_session.sparkContext.parallelize(c, numSlices)) + return self._spark_session.sparkContext.parallelize(c, numSlices) - def _get_full_path(self, relative_path): + def read_table(self, ident, options=None): + """Read a registered CSV table by ``TableIdentifier`` (public catalog access).""" + return self._catalog.read_csv(ident, options) + + def full_path(self, relative_path): + """Absolute local path for a path relative to the catalog base directory.""" return self._catalog.full_path(relative_path) + + def save_table( + self, + ident, + dataframe, + mode=None, + header=True, + replace_where=None, + csv_options=None, + overwrite_schema=False, + merge_schema=False, + ): + """Persist a DataFrame as a CSV-backed table (public catalog access).""" + self._catalog.save_dataframe( + ident, + dataframe, + mode=mode, + header=header, + replace_where=replace_where, + csv_options=csv_options, + overwrite_schema=overwrite_schema, + merge_schema=merge_schema, + ) diff --git a/tests/conftest.py b/tests/conftest.py index f7b0d0b..242a1b8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,24 @@ import os import sys +import pytest + +# Single place ensuring `testbricks` is importable (src layout) before any +# conftest/test imports it. `pythonpath = src` in pytest.ini covers it on CI; +# this keeps local invocations from any CWD working too. +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")) + +from testbricks.dbutils import configure, dbutils + # PySpark workers on Windows need an explicit Python executable path. # On this CI/dev box `python3` is not available, but `python` is. os.environ.setdefault("PYSPARK_PYTHON", sys.executable) os.environ.setdefault("PYSPARK_DRIVER_PYTHON", sys.executable) + + +@pytest.fixture +def notebook_executor(tmp_path): + """Configured dbutils executor pointed at an isolated temp directory.""" + configure(str(tmp_path), source_dir=str(tmp_path)) + yield dbutils.executor + dbutils.widgets.removeAll() diff --git a/tests/e2e_workflow/test_e2e_workflow.py b/tests/e2e_workflow/test_e2e_workflow.py index 0da004a..6e47878 100644 --- a/tests/e2e_workflow/test_e2e_workflow.py +++ b/tests/e2e_workflow/test_e2e_workflow.py @@ -1,16 +1,11 @@ import os import shutil -import sys - -# Ensure src is on the path so `testbricks` can be imported during pytest collection. -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) import pytest from testbricks.local_workflow_runner import LocalWorkflowRunner from testbricks.spark_proxy import SparkProxy - TEST_DIR = os.path.dirname(os.path.abspath(__file__)) DATA_DIR = os.path.join(TEST_DIR, "data") NOTEBOOKS_DIR = os.path.join(TEST_DIR, "notebooks") @@ -40,20 +35,29 @@ def test_e2e_workflow_runs_full_pipeline(e2e_env): runner.run_workflow(extra_globals={"spark": spark}) # NB1: USA filter → 3 customers (Alice, Charlie, Eve) - customers = spark.read.option("header", "true").option("inferSchema", "true") \ + customers = ( + spark.read.option("header", "true") + .option("inferSchema", "true") .table("silver.customers_enriched") + ) assert customers.count() == 3 assert {r.name for r in customers.collect()} == {"Alice", "Charlie", "Eve"} # NB2: amount >= 100 → 4 orders (101, 103, 105, 107) - orders = spark.read.option("header", "true").option("inferSchema", "true") \ + orders = ( + spark.read.option("header", "true") + .option("inferSchema", "true") .table("silver.orders_enriched") + ) assert orders.count() == 4 assert {r.order_id for r in orders.collect()} == {101, 103, 105, 107} # NB3: inner join → Alice(370.0), Eve(300.0); Charlie dropped (no qualifying orders) - summary = spark.read.option("header", "true").option("inferSchema", "true") \ + summary = ( + spark.read.option("header", "true") + .option("inferSchema", "true") .table("gold.customer_order_summary") + ) rows = summary.collect() assert len(rows) == 2 assert rows[0].name == "Alice" and float(rows[0].total_amount) == 370.0 diff --git a/tests/test_basic.py b/tests/test_basic.py index 2030927..e95308c 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,17 +1,16 @@ -import sys import os - -# Ensure src is on the path so `testbricks` can be imported during pytest collection. -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) +import sys import pytest -import shutil from testbricks.spark_proxy import SparkProxy - TEST_DIR = "tests/data" +WINUTILS_SKIP = pytest.mark.skipif( + sys.platform == "win32", reason="Native Spark file writers require Hadoop winutils on Windows" +) + @pytest.fixture(scope="session") def spark_session(): @@ -36,6 +35,7 @@ def temp_spark(tmp_path): def _make_df(spark_proxy, rows, columns): """Create a DataFrameWrapper from local data without using parallelize.""" from testbricks.data_frame_wrapper import DataFrameWrapper + spark_df = spark_proxy._spark_session.createDataFrame(rows, schema=columns) return DataFrameWrapper(spark_proxy, spark_df) @@ -55,7 +55,9 @@ def test_read_table_invalid_table_name_raises(self, spark): spark.read.table("drivers") def test_read_table_missing_file_raises(self, spark): - with pytest.raises(Exception): + from pyspark.sql.utils import AnalysisException + + with pytest.raises(AnalysisException): spark.read.option("header", "true").table("f1_data.missing_table") @@ -126,9 +128,9 @@ def test_format_partition_by_overwrite_schema_save_as_table(self, temp_spark): [("Alice", 30, "2024-01-01")], ["Name", "Age", "dt"], ) - df.write.format("delta").mode("overwrite").option( - "overwriteSchema", "true" - ).partitionBy("dt").saveAsTable("silver.people") + df.write.format("delta").mode("overwrite").option("overwriteSchema", "true").partitionBy( + "dt" + ).saveAsTable("silver.people") csv_path = os.path.join(temp_spark._base_path, "silver", "people.csv") assert os.path.exists(csv_path) @@ -143,9 +145,7 @@ def test_read_format_delta_table(self, spark): def test_save_as_table_three_part_name(self, temp_spark): df = temp_spark.createDataFrame([("Bob", 25)], ["Name", "Age"]) df.write.mode("overwrite").saveAsTable("main.default.people") - assert os.path.exists( - os.path.join(temp_spark._base_path, "default", "people.csv") - ) + assert os.path.exists(os.path.join(temp_spark._base_path, "default", "people.csv")) assert temp_spark.sql("SELECT * FROM default.people").count() == 1 @@ -397,9 +397,7 @@ def test_save_as_table_date_format(self, temp_spark): from datetime import date df = temp_spark.createDataFrame([(date(2026, 9, 1),)], ["dt"]) - df.write.mode("overwrite").option("dateFormat", "dd/MM/yyyy").saveAsTable( - "default.dates" - ) + df.write.mode("overwrite").option("dateFormat", "dd/MM/yyyy").saveAsTable("default.dates") csv_path = os.path.join(temp_spark._base_path, "default", "dates.csv") with open(csv_path, encoding="utf-8") as handle: contents = handle.read() @@ -408,7 +406,7 @@ def test_save_as_table_date_format(self, temp_spark): result = temp_spark.read.table("default.dates") assert result.count() == 1 - @pytest.mark.skipif(sys.platform == "win32", reason="Native Spark CSV writer requires Hadoop winutils on Windows") + @WINUTILS_SKIP def test_csv_path_write_honors_delimiter(self, temp_spark): df = _make_df(temp_spark, [(1, "a")], ["id", "name"]) df.write.mode("overwrite").option("delimiter", "|").option("header", "true").csv( @@ -422,7 +420,7 @@ def test_csv_path_write_honors_delimiter(self, temp_spark): class TestFileWriteDispatch: - @pytest.mark.skipif(sys.platform == "win32", reason="Native Spark file writers require Hadoop winutils on Windows") + @WINUTILS_SKIP def test_parquet_write_is_readable(self, temp_spark): df = _make_df(temp_spark, [("Alice", 30)], ["Name", "Age"]) df.write.mode("overwrite").parquet("output/people_parquet") @@ -431,7 +429,7 @@ def test_parquet_write_is_readable(self, temp_spark): assert result.count() == 1 assert result.collect()[0].Name == "Alice" - @pytest.mark.skipif(sys.platform == "win32", reason="Native Spark file writers require Hadoop winutils on Windows") + @WINUTILS_SKIP def test_json_write_is_readable(self, temp_spark): df = _make_df(temp_spark, [("Alice", 30)], ["Name", "Age"]) df.write.mode("overwrite").json("output/people_json") @@ -440,16 +438,14 @@ def test_json_write_is_readable(self, temp_spark): assert result.count() == 1 assert result.collect()[0].Name == "Alice" - @pytest.mark.skipif(sys.platform == "win32", reason="Native Spark file writers require Hadoop winutils on Windows") + @WINUTILS_SKIP def test_format_delta_save_writes_parquet(self, temp_spark): df = _make_df(temp_spark, [("Alice", 30)], ["Name", "Age"]) df.write.format("delta").mode("overwrite").save("output/people_delta") path = os.path.join(temp_spark._base_path, "output", "people_delta") result = temp_spark._spark_session.read.parquet(path) assert result.count() == 1 - parquet_files = [ - name for name in os.listdir(path) if name.endswith(".parquet") - ] + parquet_files = [name for name in os.listdir(path) if name.endswith(".parquet")] assert parquet_files def test_unknown_file_format_raises(self, temp_spark): @@ -486,9 +482,7 @@ def test_merge_schema_appends_missing_columns_as_nulls(self, temp_spark): first.write.mode("overwrite").saveAsTable("default.people") second = _make_df(temp_spark, [("Bob", 25, "UK")], ["Name", "Age", "Country"]) - second.write.mode("append").option("mergeSchema", "true").saveAsTable( - "default.people" - ) + second.write.mode("append").option("mergeSchema", "true").saveAsTable("default.people") result = temp_spark.sql("SELECT * FROM default.people") assert result.count() == 2 assert "Country" in result.columns @@ -504,18 +498,14 @@ def test_merge_schema_type_conflict_raises(self, temp_spark): second = temp_spark.createDataFrame([("Bob", "thirty")], ["Name", "Age"]) with pytest.raises(SchemaMismatchError, match="schema mismatch"): - second.write.mode("append").option("mergeSchema", "true").saveAsTable( - "default.people" - ) + second.write.mode("append").option("mergeSchema", "true").saveAsTable("default.people") class TestWriteTransformedTable: def test_write_transformed_table_creates_expected_csv(self, spark): # Uses the shared spark fixture because the source table lives in tests/data. df = spark.read.option("header", "true").table("f1_data.drivers") - uk = df.filter("Country = 'United Kingdom'") \ - .select("Abbreviation") \ - .distinct() + uk = df.filter("Country = 'United Kingdom'").select("Abbreviation").distinct() uk.write.mode("overwrite").saveAsTable("f1_data.uk_drivers") csv_path = os.path.join(spark._base_path, "f1_data", "uk_drivers.csv") @@ -526,7 +516,7 @@ def test_write_transformed_table_creates_expected_csv(self, spark): assert set(result.columns) == {"Abbreviation"} assert set(row.Abbreviation for row in result.collect()) == {"NOR", "RUS", "HAM", "BEA"} - @pytest.mark.skipif(sys.platform == "win32", reason="Native Spark CSV writer requires Hadoop winutils on Windows") + @WINUTILS_SKIP def test_write_csv_with_mode_and_options(self, temp_spark): df = _make_df(temp_spark, [(1,), (2,)], ["id"]) df.write.mode("overwrite").option("header", "false").csv("output/no_header") @@ -547,8 +537,8 @@ def test_option_and_options_chain(self, temp_spark): class TestDataFrameWriter: def test_writer_mode_option_chain(self, temp_spark): df = _make_df(temp_spark, [(1,)], ["id"]) - writer = df.write.format("delta").mode("overwrite").partitionBy("id").option( - "header", "true" + writer = ( + df.write.format("delta").mode("overwrite").partitionBy("id").option("header", "true") ) assert writer._mode == "overwrite" assert writer._format == "delta" diff --git a/tests/test_catalog.py b/tests/test_catalog.py index fcac9a6..be8f190 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -1,8 +1,5 @@ -import sys import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pytest from testbricks.catalog import ( @@ -106,10 +103,7 @@ class TestSqlRewrite: def test_rewrites_from_and_join(self): from testbricks.catalog import rewrite_from_join_identifiers - query = ( - "SELECT * FROM bronze.customers c " - "JOIN silver.orders o ON c.id = o.customer_id" - ) + query = "SELECT * FROM bronze.customers c JOIN silver.orders o ON c.id = o.customer_id" rewritten = rewrite_from_join_identifiers(query) assert "FROM bronze_customers c" in rewritten assert "JOIN silver_orders o" in rewritten @@ -166,9 +160,6 @@ def test_table_exists_list_tables_and_databases(self, tmp_path): filtered = spark_proxy.catalog.listTables(pattern="peo*") assert [table.name for table in filtered] == ["people"] - assert [db.name for db in spark_proxy.catalog.listDatabases(pattern="def*")] == [ - "default" - ] + assert [db.name for db in spark_proxy.catalog.listDatabases(pattern="def*")] == ["default"] ident = TableIdentifier.parse("default.people") assert spark_proxy.catalog.exists(ident) - diff --git a/tests/test_dbutils_fs.py b/tests/test_dbutils_fs.py index 8f8b612..2b3ff86 100644 --- a/tests/test_dbutils_fs.py +++ b/tests/test_dbutils_fs.py @@ -1,8 +1,5 @@ -import sys import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pytest from testbricks.dbutils import DbutilsError, configure, dbutils @@ -237,8 +234,7 @@ def test_workflow_runner_injects_dbutils(self, tmp_path): marker = store_dir / "created_by_dbutils" (source_dir / "main.py").write_text( - 'dbutils.fs.mkdirs("dbfs:/workflow_dir")\n' - f'open(r"{marker}", "w").write("ok")\n', + f'dbutils.fs.mkdirs("dbfs:/workflow_dir")\nopen(r"{marker}", "w").write("ok")\n', encoding="utf-8", ) diff --git a/tests/test_dbutils_jobs.py b/tests/test_dbutils_jobs.py index 393d104..d39956d 100644 --- a/tests/test_dbutils_jobs.py +++ b/tests/test_dbutils_jobs.py @@ -1,8 +1,5 @@ -import sys import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pytest from testbricks.dbutils import DbutilsError, configure, dbutils @@ -21,16 +18,12 @@ class TestTaskValuesSetGet: def test_set_and_get_within_current_task(self): with dbutils.jobs.taskValues.current_task("producer"): dbutils.jobs.taskValues.set(key="region", value="eu") - assert ( - dbutils.jobs.taskValues.get(taskKey="producer", key="region") == "eu" - ) + assert dbutils.jobs.taskValues.get(taskKey="producer", key="region") == "eu" def test_get_task_key_alias(self): with dbutils.jobs.taskValues.current_task("producer"): dbutils.jobs.taskValues.set(key="region", value="eu") - assert ( - dbutils.jobs.taskValues.get(task_key="producer", key="region") == "eu" - ) + assert dbutils.jobs.taskValues.get(task_key="producer", key="region") == "eu" def test_non_string_value_is_stringified(self): with dbutils.jobs.taskValues.current_task("producer"): @@ -45,17 +38,13 @@ def test_missing_key_raises_in_job(self): def test_default_used_when_key_missing_in_job(self): with dbutils.jobs.taskValues.current_task("producer"): assert ( - dbutils.jobs.taskValues.get( - taskKey="producer", key="missing", default="fallback" - ) + dbutils.jobs.taskValues.get(taskKey="producer", key="missing", default="fallback") == "fallback" ) def test_debug_value_used_outside_job(self): assert ( - dbutils.jobs.taskValues.get( - taskKey="producer", key="region", debugValue="local" - ) + dbutils.jobs.taskValues.get(taskKey="producer", key="region", debugValue="local") == "local" ) @@ -97,8 +86,6 @@ def test_isolated_notebook_run_commits_on_return(self, tmp_path): ) parent = tmp_path / "parent.py" store = dbutils.jobs.taskValues - with store.current_task("parent"), dbutils.executor.caller_context( - str(parent) - ): + with store.current_task("parent"), dbutils.executor.caller_context(str(parent)): dbutils.notebook.run("./child", 60) assert store.get(taskKey="parent", key="from_child") == "ok" diff --git a/tests/test_dbutils_library_data.py b/tests/test_dbutils_library_data.py index c1d1c94..c5e114f 100644 --- a/tests/test_dbutils_library_data.py +++ b/tests/test_dbutils_library_data.py @@ -1,8 +1,3 @@ -import sys -import os - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pandas as pd from testbricks.dbutils import dbutils diff --git a/tests/test_dbutils_notebook.py b/tests/test_dbutils_notebook.py index e6a4fc2..7bcef3e 100644 --- a/tests/test_dbutils_notebook.py +++ b/tests/test_dbutils_notebook.py @@ -1,8 +1,5 @@ -import sys import os -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pytest from testbricks.dbutils import DbutilsError, configure, dbutils @@ -33,9 +30,7 @@ def test_exit_converts_non_string_values(self): class TestNotebookRun: def test_run_returns_exit_value(self, tmp_path): child = tmp_path / "child.py" - child.write_text( - 'dbutils.notebook.exit("result")\n', encoding="utf-8" - ) + child.write_text('dbutils.notebook.exit("result")\n', encoding="utf-8") parent = tmp_path / "parent.py" with dbutils.executor.caller_context(str(parent)): @@ -67,21 +62,19 @@ def test_run_isolated_namespace(self, tmp_path): def test_run_passes_arguments_via_env(self, tmp_path): child = tmp_path / "child.py" - child.write_text( - 'import os\nRESULT = os.environ["env"]\n', encoding="utf-8" - ) + child.write_text('import os\nRESULT = os.environ["env"]\n', encoding="utf-8") parent = tmp_path / "parent.py" with dbutils.executor.caller_context(str(parent)): dbutils.notebook.run("./child", 60, {"env": "prod"}) - assert os.environ["env"] == "prod" + # Arguments reach the child via env but do not leak into the parent process. + assert "env" not in os.environ def test_run_arguments_override_widget_default(self, tmp_path): child = tmp_path / "child.py" child.write_text( - 'dbutils.widgets.text("env", "dev")\n' - 'RESULT = dbutils.widgets.get("env")\n', + 'dbutils.widgets.text("env", "dev")\nRESULT = dbutils.widgets.get("env")\n', encoding="utf-8", ) parent = tmp_path / "parent.py" @@ -89,14 +82,13 @@ def test_run_arguments_override_widget_default(self, tmp_path): with dbutils.executor.caller_context(str(parent)): dbutils.notebook.run("./child", 60, {"env": "prod"}) - assert os.environ["env"] == "prod" + # The run-time argument won over the widget default, and is cleaned up after. + assert "env" not in os.environ def test_run_relative_path(self, tmp_path): helpers = tmp_path / "helpers" helpers.mkdir() - (helpers / "setup.py").write_text( - 'dbutils.notebook.exit("ok")\n', encoding="utf-8" - ) + (helpers / "setup.py").write_text('dbutils.notebook.exit("ok")\n', encoding="utf-8") parent = tmp_path / "parent.py" with dbutils.executor.caller_context(str(parent)): diff --git a/tests/test_dbutils_secrets.py b/tests/test_dbutils_secrets.py index f021786..0bf6569 100644 --- a/tests/test_dbutils_secrets.py +++ b/tests/test_dbutils_secrets.py @@ -1,7 +1,4 @@ import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import pytest diff --git a/tests/test_dbutils_widgets.py b/tests/test_dbutils_widgets.py index f9a5a66..7e330c8 100644 --- a/tests/test_dbutils_widgets.py +++ b/tests/test_dbutils_widgets.py @@ -1,11 +1,8 @@ import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import pytest -from testbricks.dbutils import DbutilsError, configure, dbutils +from testbricks.dbutils import DbutilsError, dbutils @pytest.fixture(autouse=True) diff --git a/tests/test_fs_magic.py b/tests/test_fs_magic.py index 50b3f5f..dca6faf 100644 --- a/tests/test_fs_magic.py +++ b/tests/test_fs_magic.py @@ -1,11 +1,5 @@ -import sys -import os - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - import pytest -from testbricks.dbutils import configure, dbutils from testbricks.notebook_executor import ( transform_fs_commands, transform_run_commands, @@ -50,13 +44,6 @@ def test_run_and_sh_transforms_still_apply(self): assert "dbutils.fs.ls('dbfs:/')" in transformed -@pytest.fixture -def notebook_executor(tmp_path): - configure(str(tmp_path), source_dir=str(tmp_path)) - yield dbutils.executor - dbutils.widgets.removeAll() - - class TestExecFileFsMagic: def test_executes_fs_ls_and_continues(self, tmp_path, notebook_executor): (tmp_path / "listed").mkdir() diff --git a/tests/test_local_workflow_runner.py b/tests/test_local_workflow_runner.py index 044b7c5..e2f74ae 100644 --- a/tests/test_local_workflow_runner.py +++ b/tests/test_local_workflow_runner.py @@ -1,15 +1,11 @@ -import sys +import json import os -# Ensure src is on the path so `testbricks` can be imported during pytest collection. -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - -import json import pytest +from testbricks.dbutils import DbutilsError from testbricks.local_workflow_runner import LocalWorkflowRunner, transform_run_commands - ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) WORKFLOW_SAMPLE_PATH = os.path.join(ROOT_DIR, "tests", "data", "workflow_sample.json") DEFAULT_BASE_PATH = os.path.join(ROOT_DIR, "tests", "data") @@ -136,8 +132,7 @@ def test_base_parameters_seed_env_when_unset(self, tmp_path): source_dir.mkdir() marker = tmp_path / "seeded.txt" (source_dir / "main.py").write_text( - 'import os\n' - f'with open(r"{marker}", "w") as f: f.write(os.environ["mode"])\n', + f'import os\nwith open(r"{marker}", "w") as f: f.write(os.environ["mode"])\n', encoding="utf-8", ) @@ -165,8 +160,7 @@ def test_base_parameters_do_not_clobber_existing_env(self, tmp_path): source_dir.mkdir() marker = tmp_path / "preserved.txt" (source_dir / "main.py").write_text( - 'import os\n' - f'with open(r"{marker}", "w") as f: f.write(os.environ["mode"])\n', + f'import os\nwith open(r"{marker}", "w") as f: f.write(os.environ["mode"])\n', encoding="utf-8", ) @@ -186,9 +180,7 @@ def test_base_parameters_do_not_clobber_existing_env(self, tmp_path): os.environ["mode"] = "from_test" try: - runner = LocalWorkflowRunner( - str(source_dir), str(workflow_path), str(tmp_path) - ) + runner = LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) runner.run_workflow() finally: os.environ.pop("mode", None) @@ -299,18 +291,19 @@ def _write_notebooks(source_dir, mapping): (source_dir / f"{name}.py").write_text(body, encoding="utf-8") -class TestRunIfConditions: - def _runner(self, tmp_path, tasks, notebooks): - source_dir = tmp_path / "local_src" - source_dir.mkdir() - _write_notebooks(source_dir, notebooks) - workflow_path = tmp_path / "workflow.json" - workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") - return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) +def _runner(tmp_path, tasks, notebooks): + source_dir = tmp_path / "local_src" + source_dir.mkdir() + _write_notebooks(source_dir, notebooks) + workflow_path = tmp_path / "workflow.json" + workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") + return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) + +class TestRunIfConditions: def test_all_success_skips_when_dependency_fails(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -337,7 +330,7 @@ def test_all_success_skips_when_dependency_fails(self, tmp_path): def test_all_failed_runs_when_dependency_fails(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -364,7 +357,7 @@ def test_all_failed_runs_when_dependency_fails(self, tmp_path): def test_all_failed_skips_when_dependency_succeeds(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -390,7 +383,7 @@ def test_all_failed_skips_when_dependency_succeeds(self, tmp_path): def test_all_done_runs_after_failed_dependency(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -415,7 +408,7 @@ def test_all_done_runs_after_failed_dependency(self, tmp_path): assert log_file.read_text(encoding="utf-8").splitlines() == ["always"] def test_none_failed_skips_when_dependency_fails(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -436,7 +429,7 @@ def test_none_failed_skips_when_dependency_fails(self, tmp_path): assert runner.task_statuses["next"] == "SKIPPED" def test_none_failed_runs_when_dependency_skipped(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -469,7 +462,7 @@ def test_none_failed_runs_when_dependency_skipped(self, tmp_path): def test_at_least_one_success_runs_if_any_dep_succeeded(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -502,7 +495,7 @@ def test_at_least_one_success_runs_if_any_dep_succeeded(self, tmp_path): assert log_file.read_text(encoding="utf-8").splitlines() == ["join"] def test_depends_on_outcome_skips_when_status_does_not_match(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -537,17 +530,9 @@ def test_unknown_run_if_raises(self, tmp_path): class TestConditionTasks: - def _runner(self, tmp_path, tasks, notebooks): - source_dir = tmp_path / "local_src" - source_dir.mkdir() - _write_notebooks(source_dir, notebooks) - workflow_path = tmp_path / "workflow.json" - workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") - return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) - def test_true_branch_runs_and_false_branch_skipped(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -589,7 +574,7 @@ def test_true_branch_runs_and_false_branch_skipped(self, tmp_path): def test_false_branch_runs_when_condition_fails(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -630,7 +615,7 @@ def test_false_branch_runs_when_condition_fails(self, tmp_path): def test_greater_than_compares_numerically(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -662,7 +647,7 @@ def test_greater_than_compares_numerically(self, tmp_path): assert log_file.read_text(encoding="utf-8").splitlines() == ["true"] def test_condition_task_without_notebook_file(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -689,17 +674,9 @@ def test_missing_task_type_still_raises(self, tmp_path): class TestForEachTasks: - def _runner(self, tmp_path, tasks, notebooks): - source_dir = tmp_path / "local_src" - source_dir.mkdir() - _write_notebooks(source_dir, notebooks) - workflow_path = tmp_path / "workflow.json" - workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") - return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) - def test_runs_nested_notebook_once_per_input(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -730,7 +707,7 @@ def test_runs_nested_notebook_once_per_input(self, tmp_path): def test_inputs_from_task_values_json_list(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -765,7 +742,7 @@ def test_inputs_from_task_values_json_list(self, tmp_path): def test_child_failure_fails_for_each_task(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -799,7 +776,7 @@ def test_child_failure_fails_for_each_task(self, tmp_path): def test_literal_list_inputs(self, tmp_path): log_file = tmp_path / "execution.log" - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -828,14 +805,6 @@ def test_literal_list_inputs(self, tmp_path): class TestRepairAndRerun: - def _runner(self, tmp_path, tasks, notebooks): - source_dir = tmp_path / "local_src" - source_dir.mkdir() - _write_notebooks(source_dir, notebooks) - workflow_path = tmp_path / "workflow.json" - workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") - return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) - def test_only_runs_selected_tasks(self, tmp_path): log_file = tmp_path / "execution.log" tasks = [ @@ -858,7 +827,7 @@ def test_only_runs_selected_tasks(self, tmp_path): name: f'with open(r"{log_file}", "a") as f: f.write("{name}\\n")\n' for name in ["first", "second", "third"] } - runner = self._runner(tmp_path, tasks, notebooks) + runner = _runner(tmp_path, tasks, notebooks) runner.run_workflow(only=["second_task"]) assert log_file.read_text(encoding="utf-8").splitlines() == ["second"] assert runner.task_statuses["second_task"] == "SUCCESS" @@ -885,12 +854,12 @@ def test_from_task_runs_subgraph(self, tmp_path): name: f'with open(r"{log_file}", "a") as f: f.write("{name}\\n")\n' for name in ["first", "second", "third"] } - runner = self._runner(tmp_path, tasks, notebooks) + runner = _runner(tmp_path, tasks, notebooks) runner.run_workflow(from_task="second_task") assert log_file.read_text(encoding="utf-8").splitlines() == ["second", "third"] def test_unknown_only_task_raises(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -904,7 +873,7 @@ def test_unknown_only_task_raises(self, tmp_path): runner.run_workflow(only=["missing"]) def test_only_and_from_task_are_exclusive(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -919,18 +888,10 @@ def test_only_and_from_task_are_exclusive(self, tmp_path): class TestRetriesAndTimeouts: - def _runner(self, tmp_path, tasks, notebooks): - source_dir = tmp_path / "local_src" - source_dir.mkdir() - _write_notebooks(source_dir, notebooks) - workflow_path = tmp_path / "workflow.json" - workflow_path.write_text(json.dumps({"tasks": tasks}), encoding="utf-8") - return LocalWorkflowRunner(str(source_dir), str(workflow_path), str(tmp_path)) - def test_retries_until_success(self, tmp_path): counter = tmp_path / "counter.txt" counter.write_text("0", encoding="utf-8") - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -956,7 +917,7 @@ def test_retries_until_success(self, tmp_path): assert counter.read_text(encoding="utf-8") == "3" def test_retry_exhaustion_reraises(self, tmp_path): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -977,7 +938,7 @@ def test_retry_interval_sleeps(self, tmp_path, monkeypatch): "testbricks.local_workflow_runner.time.sleep", lambda seconds: slept.append(seconds), ) - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -994,7 +955,7 @@ def test_retry_interval_sleeps(self, tmp_path, monkeypatch): assert slept == [0.05] def test_timeout_seconds_accepted_not_enforced(self, tmp_path, capsys): - runner = self._runner( + runner = _runner( tmp_path, [ { @@ -1063,9 +1024,7 @@ def test_missing_notebook_task_raises(self, tmp_path): ) def test_invalid_notebook_path_raises(self, tmp_path, notebook_path): workflow = { - "tasks": [ - {"task_key": "t1", "notebook_task": {"notebook_path": notebook_path}} - ] + "tasks": [{"task_key": "t1", "notebook_task": {"notebook_path": notebook_path}}] } workflow_path = tmp_path / "workflow.json" workflow_path.write_text(json.dumps(workflow), encoding="utf-8") @@ -1166,11 +1125,7 @@ def test_cyclic_workflow_raises(self, tmp_path): LocalWorkflowRunner(str(tmp_path), str(workflow_path), str(tmp_path)) def test_missing_notebook_file_raises(self, tmp_path): - workflow = { - "tasks": [ - {"task_key": "t1", "notebook_task": {"notebook_path": "/a/missing"}} - ] - } + workflow = {"tasks": [{"task_key": "t1", "notebook_task": {"notebook_path": "/a/missing"}}]} workflow_path = tmp_path / "workflow.json" workflow_path.write_text(json.dumps(workflow), encoding="utf-8") @@ -1213,15 +1168,13 @@ def test_transform_magic_run_command_comment(self, tmp_path): def test_run_notebook_executes_percent_run_target(self, tmp_path): helpers_dir = tmp_path / "helpers" helpers_dir.mkdir() - (helpers_dir / "setup.py").write_text( - "SHARED_VALUE = 'from_setup'\n", encoding="utf-8" - ) + (helpers_dir / "setup.py").write_text("SHARED_VALUE = 'from_setup'\n", encoding="utf-8") main_path = tmp_path / "main.py" - main_path.write_text( - "# %run ./helpers/setup\nRESULT = SHARED_VALUE\n", encoding="utf-8" - ) + main_path.write_text("# %run ./helpers/setup\nRESULT = SHARED_VALUE\n", encoding="utf-8") - runner = LocalWorkflowRunner(str(tmp_path), _write_single_task_workflow(tmp_path), str(tmp_path)) + runner = LocalWorkflowRunner( + str(tmp_path), _write_single_task_workflow(tmp_path), str(tmp_path) + ) namespace = _notebook_namespace(str(main_path), runner) runner._execfile(str(main_path), namespace, namespace) @@ -1236,11 +1189,11 @@ def test_nested_percent_run_commands(self, tmp_path): "# %run ./common/base\nMIDDLE_VALUE = BASE_VALUE + 1\n", encoding="utf-8" ) main_path = tmp_path / "main.py" - main_path.write_text( - "# %run ./middle\nRESULT = MIDDLE_VALUE + 1\n", encoding="utf-8" - ) + main_path.write_text("# %run ./middle\nRESULT = MIDDLE_VALUE + 1\n", encoding="utf-8") - runner = LocalWorkflowRunner(str(tmp_path), _write_single_task_workflow(tmp_path), str(tmp_path)) + runner = LocalWorkflowRunner( + str(tmp_path), _write_single_task_workflow(tmp_path), str(tmp_path) + ) namespace = _notebook_namespace(str(main_path), runner) runner._execfile(str(main_path), namespace, namespace) @@ -1250,7 +1203,7 @@ def test_nested_percent_run_commands(self, tmp_path): def test_empty_percent_run_path_raises(self, tmp_path): notebook_path = str(tmp_path / "main.py") - with pytest.raises(ValueError, match="Empty %run path"): + with pytest.raises(DbutilsError, match="Empty %run path"): transform_run_commands("# %run \n", notebook_path) @@ -1259,8 +1212,7 @@ def test_top_level_exit_stops_file(self, tmp_path): notebook_path = tmp_path / "main.py" marker = tmp_path / "marker.txt" notebook_path.write_text( - 'dbutils.notebook.exit("early")\n' - f'open(r"{marker}", "w").write("ran")\n', + f'dbutils.notebook.exit("early")\nopen(r"{marker}", "w").write("ran")\n', encoding="utf-8", ) diff --git a/tests/test_sh_magic.py b/tests/test_sh_magic.py index 678f39b..f5ca909 100644 --- a/tests/test_sh_magic.py +++ b/tests/test_sh_magic.py @@ -1,13 +1,7 @@ -import sys -import os - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) +from unittest.mock import patch import pytest -from unittest.mock import patch - -from testbricks.dbutils import configure, dbutils from testbricks.notebook_exceptions import NotebookExit, ShellCommandError from testbricks.notebook_executor import ( run_shell, @@ -117,13 +111,6 @@ def test_missing_bash_raises(self): run_shell("echo hi") -@pytest.fixture -def notebook_executor(tmp_path): - configure(str(tmp_path), source_dir=str(tmp_path)) - yield dbutils.executor - dbutils.widgets.removeAll() - - class TestExecFileShMagic: def test_executes_sh_and_continues(self, tmp_path, notebook_executor, capsys): notebook = tmp_path / "main.py" @@ -153,9 +140,7 @@ def test_percent_run_child_can_use_sh(self, tmp_path, notebook_executor, capsys) main = tmp_path / "main.py" main.write_text("# %run ./child\nAFTER = True\n", encoding="utf-8") namespace = {"__name__": "__main__", "__file__": str(main)} - namespace["__run_notebook__"] = lambda path: notebook_executor.run_shared( - path, namespace - ) + namespace["__run_notebook__"] = lambda path: notebook_executor.run_shared(path, namespace) notebook_executor.exec_file(str(main), namespace, top_level=True) captured = capsys.readouterr() assert "from_child" in captured.out @@ -167,4 +152,3 @@ def test_isolated_run_propagates_shell_error(self, tmp_path, notebook_executor): with notebook_executor.caller_context(str(parent)): with pytest.raises(ShellCommandError): notebook_executor.run_isolated("./child") -