Skip to content

Commit

Permalink
Added content assertion for assert_file_uploaded and `assert_file_d…
Browse files Browse the repository at this point in the history
…bfs_uploaded` in `MockInstallation` (databrickslabs#101)

Add `expected: bytes` parameter to assert that the content uploaded to
MockInstallation is correct
  • Loading branch information
nkvuong authored and aminmovahed-db committed May 8, 2024
1 parent cee03b4 commit 90ff5e1
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
21 changes: 14 additions & 7 deletions src/databricks/labs/blueprint/installation.py
Original file line number Diff line number Diff line change
Expand Up @@ -903,20 +903,27 @@ 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 | 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):
self._assert_upload(filename, self._dbfs)
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)

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 | None = None):
if isinstance(filename, re.Pattern):
for name in loc.keys():
if filename.match(name):
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:
assert loc[filename] == expected, f"{filename} content missmatch"
6 changes: 6 additions & 0 deletions tests/unit/test_installation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit 90ff5e1

Please sign in to comment.