From 19b5d46749bf489d8a43fa6b65d27058b70eaf93 Mon Sep 17 00:00:00 2001 From: Vuong Date: Tue, 7 May 2024 12:06:34 +0100 Subject: [PATCH 1/3] add content assertion --- src/databricks/labs/blueprint/installation.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index b1e469b..dfeb54f 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -903,20 +903,24 @@ def assert_file_written(self, filename: str, expected: Any): actual = self._overwrites[filename] assert expected == actual, f"{filename} content missmatch" - def assert_file_uploaded(self, filename): - self._assert_upload(filename, self._uploads) + def assert_file_uploaded(self, filename, expected: bytes): + """Asserts that a file was uploaded with the expected content""" + self._assert_upload(filename, self._uploads, expected) - def assert_file_dbfs_uploaded(self, filename): - self._assert_upload(filename, self._dbfs) + def assert_file_dbfs_uploaded(self, filename, expected: bytes): + """Asserts that a file was uploaded to DBFS with the expected content""" + self._assert_upload(filename, self._dbfs, expected) def assert_removed(self): assert self._removed @staticmethod - def _assert_upload(filename: Any, loc: dict[str, bytes]): + def _assert_upload(filename: Any, loc: dict[str, bytes], expected: bytes): if isinstance(filename, re.Pattern): for name in loc.keys(): if filename.match(name): + assert loc[name] == expected, f"{filename} content missmatch" return raise AssertionError(f'Cannot find {filename.pattern} among {", ".join(loc.keys())}') assert filename in loc, f"{filename} had no writes" + assert loc[filename] == expected, f"{filename} content missmatch" From 03e00c5b4bb666e8495a6c7bc23097926ad01b6a Mon Sep 17 00:00:00 2001 From: Vuong Date: Tue, 7 May 2024 12:11:26 +0100 Subject: [PATCH 2/3] optional --- src/databricks/labs/blueprint/installation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index dfeb54f..b94d591 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -903,11 +903,11 @@ def assert_file_written(self, filename: str, expected: Any): actual = self._overwrites[filename] assert expected == actual, f"{filename} content missmatch" - def assert_file_uploaded(self, filename, expected: bytes): + def assert_file_uploaded(self, filename, expected: bytes | None = None): """Asserts that a file was uploaded with the expected content""" self._assert_upload(filename, self._uploads, expected) - def assert_file_dbfs_uploaded(self, filename, expected: bytes): + def assert_file_dbfs_uploaded(self, filename, expected: bytes | None = None): """Asserts that a file was uploaded to DBFS with the expected content""" self._assert_upload(filename, self._dbfs, expected) @@ -915,12 +915,14 @@ def assert_removed(self): assert self._removed @staticmethod - def _assert_upload(filename: Any, loc: dict[str, bytes], expected: bytes): + def _assert_upload(filename: Any, loc: dict[str, bytes], expected: bytes | None = None): if isinstance(filename, re.Pattern): for name in loc.keys(): if filename.match(name): - assert loc[name] == expected, f"{filename} content missmatch" + if expected: + assert loc[name] == expected, f"{filename} content missmatch" return raise AssertionError(f'Cannot find {filename.pattern} among {", ".join(loc.keys())}') assert filename in loc, f"{filename} had no writes" - assert loc[filename] == expected, f"{filename} content missmatch" + if expected: + assert loc[filename] == expected, f"{filename} content missmatch" From 2144b171d134f69dbf981c163633ea917593a4ab Mon Sep 17 00:00:00 2001 From: Vuong Date: Tue, 7 May 2024 12:17:53 +0100 Subject: [PATCH 3/3] add test --- src/databricks/labs/blueprint/installation.py | 9 +++++---- tests/unit/test_installation.py | 6 ++++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index b94d591..71db98f 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -918,10 +918,11 @@ def assert_removed(self): def _assert_upload(filename: Any, loc: dict[str, bytes], expected: bytes | None = None): if isinstance(filename, re.Pattern): for name in loc.keys(): - if filename.match(name): - if expected: - assert loc[name] == expected, f"{filename} content missmatch" - return + if not filename.match(name): + continue + if expected: + assert loc[name] == expected, f"{filename} content missmatch" + return raise AssertionError(f'Cannot find {filename.pattern} among {", ".join(loc.keys())}') assert filename in loc, f"{filename} had no writes" if expected: diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index c703db2..6b67219 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -445,3 +445,9 @@ def test_load_empty_data_class(): ) load = installation.load(ComplexClass, filename="backups/complex-class.json") assert load == complex_class + + +def test_assert_file_uploaded(): + installation = MockInstallation() + installation.upload("foo", b"bar") + installation.assert_file_uploaded("foo", b"bar")