common_test.go 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. package protocol
  2. import (
  3. "io"
  4. "time"
  5. )
  6. type TestModel struct {
  7. data []byte
  8. repo string
  9. name string
  10. offset int64
  11. size int
  12. closedCh chan bool
  13. }
  14. func newTestModel() *TestModel {
  15. return &TestModel{
  16. closedCh: make(chan bool),
  17. }
  18. }
  19. func (t *TestModel) Index(nodeID string, files []FileInfo) {
  20. }
  21. func (t *TestModel) IndexUpdate(nodeID string, files []FileInfo) {
  22. }
  23. func (t *TestModel) Request(nodeID, repo, name string, offset int64, size int) ([]byte, error) {
  24. t.repo = repo
  25. t.name = name
  26. t.offset = offset
  27. t.size = size
  28. return t.data, nil
  29. }
  30. func (t *TestModel) Close(nodeID string, err error) {
  31. close(t.closedCh)
  32. }
  33. func (t *TestModel) isClosed() bool {
  34. select {
  35. case <-t.closedCh:
  36. return true
  37. case <-time.After(1 * time.Second):
  38. return false // Timeout
  39. }
  40. }
  41. type ErrPipe struct {
  42. io.PipeWriter
  43. written int
  44. max int
  45. err error
  46. closed bool
  47. }
  48. func (e *ErrPipe) Write(data []byte) (int, error) {
  49. if e.closed {
  50. return 0, e.err
  51. }
  52. if e.written+len(data) > e.max {
  53. n, _ := e.PipeWriter.Write(data[:e.max-e.written])
  54. e.PipeWriter.CloseWithError(e.err)
  55. e.closed = true
  56. return n, e.err
  57. }
  58. return e.PipeWriter.Write(data)
  59. }