Example #1
0
    def test_download_dir_with_basename(self, client):
        dirname1 = tempfile.mkdtemp()
        dirname2 = tempfile.mkdtemp(prefix=dirname1 + "/")

        def mkfile(fname):
            return open(fname, "w")

        client.return_value.get_blob_to_path.side_effect = mkfile

        store = AzureBlobStoreService()
        store.set_connection(connection=client)

        blob_path = "/path/to/"
        azure_url = self.wasbs_base + blob_path
        rel_path2 = dirname2.split("/")[-1]

        # Mock return list
        obj_mock1 = BlobProperties()
        obj_mock1.size = 42
        obj_mock1.name = blob_path + "foo/test1.txt"

        obj_mock2 = BlobProperties()
        obj_mock2.size = 42
        obj_mock2.name = blob_path + "foo/test2.txt"

        subdir_mock = BlobPrefix()
        subdir_mock.name = blob_path + "foo/" + rel_path2

        obj_mock3 = BlobProperties()
        obj_mock3.size = 42
        obj_mock3.name = blob_path + "foo/" + rel_path2 + "/" + "test3.txt"

        # Create some files to return
        def list_side_effect(container_name, prefix, delimiter="/", marker=None):
            if prefix == blob_path + "foo/":
                return MockBlobList([subdir_mock, obj_mock1, obj_mock2])
            return MockBlobList([obj_mock3])

        client.return_value.walk_blobs.side_effect = list_side_effect

        dirname3 = tempfile.mkdtemp()

        # Test without basename
        store.download_dir(
            blob=azure_url + "foo", local_path=dirname3, use_basename=True
        )
        call_args_list = client.get_container_client().download_blob.call_args_list
        for call_args in call_args_list:
            call_arg1 = call_args[0]
            if call_arg1 == "{}foo/test1.txt".format(blob_path):
                assert os.path.exists("{}/foo/test1.txt".format(dirname3))
            elif call_arg1 == "{}foo/test2.txt".format(blob_path):
                assert os.path.exists("{}/foo/test2.txt".format(dirname3))
            elif call_arg1 == "{}foo/{}/test3.txt".format(blob_path, rel_path2):
                assert os.path.exists("{}/foo/{}/test3.txt".format(dirname3, rel_path2))
            else:
                assert False
Example #2
0
    def test_download_dir_with_basename(self, client):
        dirname1 = tempfile.mkdtemp()
        dirname2 = tempfile.mkdtemp(prefix=dirname1 + "/")

        def mkfile(container, cloud_path, fname):
            return open(fname, "w")

        client.return_value.get_blob_to_path.side_effect = mkfile

        store = AzureBlobStoreService()

        blob_path = "/path/to/"
        azure_url = self.wasbs_base + blob_path
        rel_path2 = dirname2.split("/")[-1]

        # Mock return list
        blob_props = BlobProperties()
        blob_props.content_length = 42
        obj_mock1 = Blob(blob_path + "foo/test1.txt", props=blob_props)

        blob_props = BlobProperties()
        blob_props.content_length = 42
        obj_mock2 = Blob(blob_path + "foo/test2.txt", props=blob_props)

        subdir_mock = BlobPrefix()
        subdir_mock.name = blob_path + "foo/" + rel_path2

        blob_props = BlobProperties()
        blob_props.content_length = 42
        obj_mock3 = Blob(
            blob_path + "foo/" + rel_path2 + "/" + "test3.txt", props=blob_props
        )

        # Create some files to return
        def list_side_effect(container_name, prefix, delimiter="/", marker=None):
            if prefix == blob_path + "foo/":
                return MockBlobList([subdir_mock, obj_mock1, obj_mock2])
            return MockBlobList([obj_mock3])

        client.return_value.list_blobs.side_effect = list_side_effect

        dirname3 = tempfile.mkdtemp()

        # Test without basename
        store.download_dir(
            blob=azure_url + "foo", local_path=dirname3, use_basename=True
        )
        client.return_value.get_blob_to_path.assert_has_calls(
            [
                mock.call(
                    "container",
                    "{}foo/test1.txt".format(blob_path),
                    "{}/foo/test1.txt".format(dirname3),
                ),
                mock.call(
                    "container",
                    "{}foo/test2.txt".format(blob_path),
                    "{}/foo/test2.txt".format(dirname3),
                ),
                mock.call(
                    "container",
                    "{}foo/{}/test3.txt".format(blob_path, rel_path2),
                    "{}/foo/{}/test3.txt".format(dirname3, rel_path2),
                ),
            ],
            any_order=True,
        )