tests/utils/cas.py: Extend CAS server helper

#79

tests/utils/cas.py

@@ -24,11 +24,13 @@ import grpc
 import psutil
 import pytest_cov
 
+from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from buildgrid.server.cas.service import ByteStreamService
 from buildgrid.server.cas.service import ContentAddressableStorageService
 from buildgrid.server.cas.instance import ByteStreamInstance
 from buildgrid.server.cas.instance import ContentAddressableStorageInstance
 from buildgrid.server.cas.storage.disk import DiskStorage
+from buildgrid.utils import create_digest, merkle_maker
 
 
 @contextmanager
@@ -103,6 +105,15 @@ class Server:
     def has(self, digest):
         return self.__storage.has_blob(digest)
 
+    def store_blob(self, blob):
+        digest = create_digest(blob)
+        write_buffer = self.__storage.begin_write(digest)
+        write_buffer.write(blob)
+
+        self.__storage.commit_write(digest, write_buffer)
+
+        return digest
+
     def compare_blobs(self, digest, blob):
         if not self.__storage.has_blob(digest):
             return False
@@ -112,6 +123,16 @@ class Server:
 
         return blob == stored_blob
 
+    def store_message(self, message):
+        message_blob = message.SerializeToString()
+        message_digest = create_digest(message_blob)
+        write_buffer = self.__storage.begin_write(message_digest)
+        write_buffer.write(message_blob)
+
+        self.__storage.commit_write(message_digest, write_buffer)
+
+        return message_digest
+
     def compare_messages(self, digest, message):
         if not self.__storage.has_blob(digest):
             return False
@@ -123,6 +144,17 @@ class Server:
 
         return message_blob == stored_blob
 
+    def store_file(self, file_path):
+        with open(file_path, 'rb') as file_bytes:
+            file_blob = file_bytes.read()
+        file_digest = create_digest(file_blob)
+        write_buffer = self.__storage.begin_write(file_digest)
+        write_buffer.write(file_blob)
+
+        self.__storage.commit_write(file_digest, write_buffer)
+
+        return file_digest
+
     def compare_files(self, digest, file_path):
         if not self.__storage.has_blob(digest):
             return False
@@ -135,6 +167,53 @@ class Server:
 
         return file_blob == stored_blob
 
+    def store_folder(self, folder_path):
+        last_digest = None
+        for blob, digest in merkle_maker(folder_path):
+            write_buffer = self.__storage.begin_write(digest)
+            write_buffer.write(blob)
+
+            self.__storage.commit_write(digest, write_buffer)
+            last_digest = digest
+
+        return last_digest
+
+    def check_local_folders(self, digest, folder_path):
+        if not self.__storage.has_blob(digest):
+            return False
+
+        def __compare_folders(digest, path):
+            directory = remote_execution_pb2.Directory()
+            directory.ParseFromString(self.__storage.get_blob(digest).read())
+
+            for file_node in directory.files:
+                file_path = os.path.join(path, file_node.name)
+
+                assert os.path.isfile(file_path)
+                assert not os.path.islink(file_path)
+                if file_node.is_executable:
+                    assert os.access(file_path, os.X_OK)
+
+                assert self.compare_files(file_node.digest, file_path)
+
+            for symlink_node in directory.symlinks:
+                symlink_path = os.path.join(path, symlink_node.name)
+
+                assert os.path.islink(symlink_path)
+                assert os.readlink(symlink_path) == symlink_node.target
+
+            for directory_node in directory.directories:
+                directory_path = os.path.join(path, directory_node.name)
+
+                assert os.path.exists(directory_path)
+                assert not os.path.islink(directory_path)
+
+                return __compare_folders(directory_node.digest, directory_path)
+
+            return True
+
+        return __compare_folders(digest, folder_path)
+
     def quit(self):
         if self.__process:
             self.__process.terminate()