github.com/mika/distribution@v2.2.2-0.20160108133430-a75790e3d8e0+incompatible/registry/storage/driver/azure/blockblob_test.go (about)

     1  package azure
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  
     9  	azure "github.com/Azure/azure-sdk-for-go/storage"
    10  )
    11  
    12  type StorageSimulator struct {
    13  	blobs map[string]*BlockBlob
    14  }
    15  
    16  type BlockBlob struct {
    17  	blocks    map[string]*DataBlock
    18  	blockList []string
    19  }
    20  
    21  type DataBlock struct {
    22  	data      []byte
    23  	committed bool
    24  }
    25  
    26  func (s *StorageSimulator) path(container, blob string) string {
    27  	return fmt.Sprintf("%s/%s", container, blob)
    28  }
    29  
    30  func (s *StorageSimulator) BlobExists(container, blob string) (bool, error) {
    31  	_, ok := s.blobs[s.path(container, blob)]
    32  	return ok, nil
    33  }
    34  
    35  func (s *StorageSimulator) GetBlob(container, blob string) (io.ReadCloser, error) {
    36  	bb, ok := s.blobs[s.path(container, blob)]
    37  	if !ok {
    38  		return nil, fmt.Errorf("blob not found")
    39  	}
    40  
    41  	var readers []io.Reader
    42  	for _, bID := range bb.blockList {
    43  		readers = append(readers, bytes.NewReader(bb.blocks[bID].data))
    44  	}
    45  	return ioutil.NopCloser(io.MultiReader(readers...)), nil
    46  }
    47  
    48  func (s *StorageSimulator) GetSectionReader(container, blob string, start, length int64) (io.ReadCloser, error) {
    49  	r, err := s.GetBlob(container, blob)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	b, err := ioutil.ReadAll(r)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	return ioutil.NopCloser(bytes.NewReader(b[start : start+length])), nil
    58  }
    59  
    60  func (s *StorageSimulator) CreateBlockBlob(container, blob string) error {
    61  	path := s.path(container, blob)
    62  	bb := &BlockBlob{
    63  		blocks:    make(map[string]*DataBlock),
    64  		blockList: []string{},
    65  	}
    66  	s.blobs[path] = bb
    67  	return nil
    68  }
    69  
    70  func (s *StorageSimulator) PutBlock(container, blob, blockID string, chunk []byte) error {
    71  	path := s.path(container, blob)
    72  	bb, ok := s.blobs[path]
    73  	if !ok {
    74  		return fmt.Errorf("blob not found")
    75  	}
    76  	data := make([]byte, len(chunk))
    77  	copy(data, chunk)
    78  	bb.blocks[blockID] = &DataBlock{data: data, committed: false} // add block to blob
    79  	return nil
    80  }
    81  
    82  func (s *StorageSimulator) GetBlockList(container, blob string, blockType azure.BlockListType) (azure.BlockListResponse, error) {
    83  	resp := azure.BlockListResponse{}
    84  	bb, ok := s.blobs[s.path(container, blob)]
    85  	if !ok {
    86  		return resp, fmt.Errorf("blob not found")
    87  	}
    88  
    89  	// Iterate committed blocks (in order)
    90  	if blockType == azure.BlockListTypeAll || blockType == azure.BlockListTypeCommitted {
    91  		for _, blockID := range bb.blockList {
    92  			b := bb.blocks[blockID]
    93  			block := azure.BlockResponse{
    94  				Name: blockID,
    95  				Size: int64(len(b.data)),
    96  			}
    97  			resp.CommittedBlocks = append(resp.CommittedBlocks, block)
    98  		}
    99  
   100  	}
   101  
   102  	// Iterate uncommitted blocks (in no order)
   103  	if blockType == azure.BlockListTypeAll || blockType == azure.BlockListTypeCommitted {
   104  		for blockID, b := range bb.blocks {
   105  			block := azure.BlockResponse{
   106  				Name: blockID,
   107  				Size: int64(len(b.data)),
   108  			}
   109  			if !b.committed {
   110  				resp.UncommittedBlocks = append(resp.UncommittedBlocks, block)
   111  			}
   112  		}
   113  	}
   114  	return resp, nil
   115  }
   116  
   117  func (s *StorageSimulator) PutBlockList(container, blob string, blocks []azure.Block) error {
   118  	bb, ok := s.blobs[s.path(container, blob)]
   119  	if !ok {
   120  		return fmt.Errorf("blob not found")
   121  	}
   122  
   123  	var blockIDs []string
   124  	for _, v := range blocks {
   125  		bl, ok := bb.blocks[v.ID]
   126  		if !ok { // check if block ID exists
   127  			return fmt.Errorf("Block id '%s' not found", v.ID)
   128  		}
   129  		bl.committed = true
   130  		blockIDs = append(blockIDs, v.ID)
   131  	}
   132  
   133  	// Mark all other blocks uncommitted
   134  	for k, b := range bb.blocks {
   135  		inList := false
   136  		for _, v := range blockIDs {
   137  			if k == v {
   138  				inList = true
   139  				break
   140  			}
   141  		}
   142  		if !inList {
   143  			b.committed = false
   144  		}
   145  	}
   146  
   147  	bb.blockList = blockIDs
   148  	return nil
   149  }
   150  
   151  func NewStorageSimulator() StorageSimulator {
   152  	return StorageSimulator{
   153  		blobs: make(map[string]*BlockBlob),
   154  	}
   155  }