Skip to content

Commit 8cf5dc4

Browse files
authored
Add method to detect rows are written to the MockBackend (#292)
See [this comment](https://github.com/databrickslabs/ucx/pull/2625/files#r1758391830) on motivation
1 parent 02ed2c8 commit 8cf5dc4

File tree

2 files changed

+49
-0
lines changed

2 files changed

+49
-0
lines changed

src/databricks/labs/lsql/backends.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,18 @@ def rows_written_for(self, full_name: str, mode: str) -> list[DataclassInstance]
372372
rows += stub_rows
373373
return rows
374374

375+
def has_rows_written_for(self, full_name: str) -> bool:
376+
"""Check if the table has been written to.
377+
378+
This method allows to differentiate between "never written to the table" (returns False) and "zero rows written
379+
to the table" (return True).
380+
Otherwise, the check is the same as: `assert (not) mock_backend.rows_written_for(full_name, mode)`
381+
"""
382+
for stub_full_name, _, _ in self._save_table:
383+
if stub_full_name == full_name:
384+
return True
385+
return False
386+
375387
@staticmethod
376388
def rows(*column_names: str):
377389
"""This method is used to create rows for the mock backend."""

tests/unit/test_backends.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,43 @@ def test_mock_backend_overwrite():
451451
]
452452

453453

454+
@pytest.mark.parametrize("mode", ["append", "overwrite"])
455+
def test_mock_backend_has_no_rows_written(mode) -> None:
456+
mock_backend = MockBackend()
457+
# There are no rows written
458+
assert not mock_backend.has_rows_written_for("a.b.c")
459+
# and the results contains no rows
460+
assert not mock_backend.rows_written_for("a.b.c", mode)
461+
462+
463+
@pytest.mark.parametrize("mode", ["append", "overwrite"])
464+
def test_mock_backend_has_zero_rows_written(mode) -> None:
465+
mock_backend = MockBackend()
466+
mock_backend.save_table("a.b.c", [], Foo, mode)
467+
# There are rows written
468+
assert mock_backend.has_rows_written_for("a.b.c")
469+
# while the results contains no rows
470+
assert not mock_backend.rows_written_for("a.b.c", mode)
471+
472+
473+
@pytest.mark.parametrize("mode", ["append", "overwrite"])
474+
def test_mock_backend_has_rows_written_for_after_first_write(mode) -> None:
475+
mock_backend = MockBackend()
476+
mock_backend.save_table("a.b.c", [Foo("a1", True), Foo("c2", False)], Foo, mode)
477+
assert mock_backend.has_rows_written_for("a.b.c")
478+
assert not mock_backend.has_rows_written_for("a.b.d")
479+
480+
481+
@pytest.mark.parametrize("first_mode", ["append", "overwrite"])
482+
@pytest.mark.parametrize("second_mode", ["append", "overwrite"])
483+
def test_mock_backend_has_rows_written_for_after_two_writes(first_mode, second_mode) -> None:
484+
mock_backend = MockBackend()
485+
mock_backend.save_table("a.b.c", [Foo("a1", True), Foo("c2", False)], Foo, first_mode)
486+
mock_backend.save_table("a.b.c", [Foo("aaa", True), Foo("bbb", False)], Foo, second_mode)
487+
assert mock_backend.has_rows_written_for("a.b.c")
488+
assert not mock_backend.has_rows_written_for("a.b.d")
489+
490+
454491
@dataclass
455492
class Nested:
456493
foo: Foo

0 commit comments

Comments
 (0)