Merge pull request #1275 from antsystem/feat/add-lru-queue-and-mockfileserver-prepare-pr

add lru queue and mockFileServer

Author: lowzj

dfget/core/helper/helper_test.go

@@ -17,12 +17,17 @@
 package helper
 
 import (
+	"context"
+	"fmt"
+	"io/ioutil"
 	"reflect"
 	"runtime"
 	"strings"
 	"testing"
 
 	"github.com/dragonflyoss/Dragonfly/dfget/config"
+	"github.com/dragonflyoss/Dragonfly/pkg/httputils"
+
 	"github.com/go-check/check"
 )
 
@@ -74,3 +79,115 @@ func (s *HelperTestSuite) TestDownloadPattern(c *check.C) {
 			check.Commentf("f:%v pattern:%s", name(v.f), v.pattern))
 	}
 }
+
+func (s *HelperTestSuite) readFromFileServer(port int, path string, off int64, size int64) ([]byte, error) {
+	url := fmt.Sprintf("http://127.0.0.1:%d/%s", port, path)
+	header := map[string]string{}
+
+	if size > 0 {
+		header["Range"] = fmt.Sprintf("bytes=%d-%d", off, off+size-1)
+	}
+
+	resp, err := httputils.HTTPGet(url, header)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode >= 400 {
+		return nil, fmt.Errorf("resp code %d", resp.StatusCode)
+	}
+
+	return ioutil.ReadAll(resp.Body)
+}
+
+func (s *HelperTestSuite) getRepeatStr(data []byte, size int64) []byte {
+	for {
+		if int64(len(data)) >= size {
+			break
+		}
+
+		newData := make([]byte, len(data)*2)
+		copy(newData, data)
+		copy(newData[len(data):], data)
+		data = newData
+	}
+
+	return data[:size]
+}
+
+func (s *HelperTestSuite) TestMockFileServer(c *check.C) {
+	var (
+		data []byte
+		err  error
+	)
+
+	// run server
+	mfs := NewMockFileServer()
+	err = mfs.StartServer(context.Background(), 10011)
+	c.Assert(err, check.IsNil)
+
+	// register fileA
+	err = mfs.RegisterFile("fileA", 100, "a")
+	c.Assert(err, check.IsNil)
+
+	// test fileA
+	// read 0-9
+	data, err = s.readFromFileServer(mfs.Port, "fileA", 0, 10)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("a"), 10)))
+
+	// read 1-20
+	data, err = s.readFromFileServer(mfs.Port, "fileA", 1, 20)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("a"), 20)))
+
+	// read 1-100
+	data, err = s.readFromFileServer(mfs.Port, "fileA", 1, 99)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("a"), 99)))
+
+	// read all
+	data, err = s.readFromFileServer(mfs.Port, "fileA", 0, -1)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("a"), 100)))
+
+	// read 0-99
+	data, err = s.readFromFileServer(mfs.Port, "fileA", 0, 100)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("a"), 100)))
+
+	// register fileB
+	err = mfs.RegisterFile("fileB", 10000, "abcde")
+	c.Assert(err, check.IsNil)
+
+	// read 0-9
+	data, err = s.readFromFileServer(mfs.Port, "fileB", 0, 10)
+	c.Assert(err, check.IsNil)
+	c.Assert(string(data), check.Equals, string(s.getRepeatStr([]byte("abcde"), 10)))
+
+	// read 8-100
+	data, err = s.readFromFileServer(mfs.Port, "fileB", 8, 93)
+	c.Assert(err, check.IsNil)
+	expectStr := s.getRepeatStr([]byte("abcde"), 101)
+	c.Assert(string(data), check.Equals, string(expectStr[8:]))
+
+	// read 1000-9999
+	data, err = s.readFromFileServer(mfs.Port, "fileB", 1000, 9000)
+	c.Assert(err, check.IsNil)
+	expectStr = s.getRepeatStr([]byte("abcde"), 10000)
+	c.Assert(string(data), check.Equals, string(expectStr[1000:]))
+
+	// read 1001-9000
+	data, err = s.readFromFileServer(mfs.Port, "fileB", 1001, 8000)
+	c.Assert(err, check.IsNil)
+	expectStr = s.getRepeatStr([]byte("abcde"), 9001)
+	c.Assert(string(data), check.Equals, string(expectStr[1001:]))
+
+	// read all
+	data, err = s.readFromFileServer(mfs.Port, "fileB", 0, -1)
+	c.Assert(err, check.IsNil)
+	expectStr = s.getRepeatStr([]byte("abcde"), 10000)
+	c.Assert(string(data), check.Equals, string(expectStr))
+}

dfget/core/helper/test_helper.go

@@ -17,19 +17,26 @@
 package helper
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math"
 	"math/rand"
+	"net"
+	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
+	"sync"
 
 	api_types "github.com/dragonflyoss/Dragonfly/apis/types"
 	"github.com/dragonflyoss/Dragonfly/dfget/config"
 	"github.com/dragonflyoss/Dragonfly/dfget/core/api"
 	"github.com/dragonflyoss/Dragonfly/dfget/types"
 	"github.com/dragonflyoss/Dragonfly/pkg/constants"
 	"github.com/dragonflyoss/Dragonfly/pkg/fileutils"
+	"github.com/dragonflyoss/Dragonfly/pkg/httputils"
 
 	"github.com/sirupsen/logrus"
 )
@@ -210,3 +217,180 @@ func CreateRegisterFunc() RegisterFuncType {
 		return nil, nil
 	}
 }
+
+// MockFileServer mocks the file server.
+type MockFileServer struct {
+	sync.Mutex
+	Port    int
+	fileMap map[string]*mockFile
+	sr      *http.Server
+}
+
+func NewMockFileServer() *MockFileServer {
+	return &MockFileServer{
+		fileMap: make(map[string]*mockFile),
+	}
+}
+
+// StartServer asynchronously starts the server, it will not be blocked.
+func (fs *MockFileServer) StartServer(ctx context.Context, port int) error {
+	addr, err := net.ResolveTCPAddr("", fmt.Sprintf(":%d", port))
+	if err != nil {
+		return err
+	}
+
+	l, err := net.ListenTCP("tcp", addr)
+	if err != nil {
+		return err
+	}
+
+	fs.Port = addr.Port
+	sr := &http.Server{}
+	sr.Handler = fs
+	fs.sr = sr
+
+	go func() {
+		if err := fs.sr.Serve(l); err != nil {
+			panic(err)
+		}
+	}()
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				fs.sr.Close()
+				return
+			}
+		}
+	}()
+
+	return nil
+}
+
+func (fs *MockFileServer) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
+	var (
+		err      error
+		reqRange *httputils.RangeStruct
+	)
+
+	if req.Method != http.MethodGet {
+		resp.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	path := req.URL.Path
+	path = strings.Trim(path, "/")
+
+	fs.Lock()
+	mf, exist := fs.fileMap[path]
+	fs.Unlock()
+
+	if !exist {
+		resp.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	rangeSt := []*httputils.RangeStruct{}
+	rangeStr := req.Header.Get("Range")
+
+	if rangeStr != "" {
+		rangeSt, err = httputils.GetRangeSE(rangeStr, math.MaxInt64)
+		if err != nil {
+			resp.WriteHeader(http.StatusBadRequest)
+			return
+		}
+	}
+
+	if len(rangeSt) > 0 {
+		reqRange = rangeSt[0]
+	}
+
+	fs.MockResp(resp, mf, reqRange)
+}
+
+func (fs *MockFileServer) RegisterFile(path string, size int64, repeatStr string) error {
+	fs.Lock()
+	defer fs.Unlock()
+
+	path = strings.Trim(path, "/")
+	_, exist := fs.fileMap[path]
+	if exist {
+		return os.ErrExist
+	}
+
+	data := []byte(repeatStr)
+	if len(data) < 1024 {
+		for {
+			newData := make([]byte, len(data)*2)
+			copy(newData, data)
+			copy(newData[len(data):], data)
+			data = newData
+
+			if len(data) >= 1024 {
+				break
+			}
+		}
+	}
+
+	fs.fileMap[path] = &mockFile{
+		path:      path,
+		size:      size,
+		repeatStr: data,
+	}
+
+	return nil
+}
+
+func (fs *MockFileServer) UnRegisterFile(path string) {
+	fs.Lock()
+	defer fs.Unlock()
+
+	delete(fs.fileMap, strings.Trim(path, "/"))
+}
+
+func (fs *MockFileServer) MockResp(resp http.ResponseWriter, mf *mockFile, rangeSt *httputils.RangeStruct) {
+	var (
+		respCode int
+		start    int64
+		end      = mf.size - 1
+	)
+	if rangeSt != nil {
+		start = rangeSt.StartIndex
+		if rangeSt.EndIndex < end {
+			end = rangeSt.EndIndex
+		}
+		respCode = http.StatusPartialContent
+	} else {
+		respCode = http.StatusOK
+	}
+
+	resp.Header().Set("Content-Length", fmt.Sprintf("%d", end-start+1))
+	resp.WriteHeader(respCode)
+
+	repeatStrLen := int64(len(mf.repeatStr))
+	strIndex := start % int64(repeatStrLen)
+
+	for {
+		if start > end {
+			break
+		}
+
+		copyDataLen := repeatStrLen - strIndex
+		if copyDataLen > (end - start + 1) {
+			copyDataLen = end - start + 1
+		}
+
+		resp.Write(mf.repeatStr[strIndex : copyDataLen+strIndex])
+		strIndex = 0
+		start += copyDataLen
+	}
+
+	return
+}
+
+type mockFile struct {
+	path      string
+	size      int64
+	repeatStr []byte
+}