func (b *Broker) fetchFile(controller Controller, message *messages.Client, conn *websocket.Conn) {
fetchFile := message.GetFetchFile()
err := controller.FetchFile(fetchFile.GetFileId(), fetchFile.GetName())
if err != nil && err.Uri == "" {
b.writeError(err.Error(), message, conn)
return
}
b.writeFileLocation(err, message, conn)
}
func (c *Connection) writeMessage(msg *messages.Client) {
data, err := msg.Marshal()
if err != nil {
c.log.Panic("Unable to marshal message", err)
}
if err = c.conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
c.writeError(err)
return
}
}
func (b *Broker) writeReadData(data []byte, message *messages.Client, conn *websocket.Conn) {
msgType := messages.Server_ReadData
server := &messages.Server{
MessageType: &msgType,
MessageId: proto.Uint64(message.GetMessageId()),
ReadData: &messages.ReadData{
Data: data,
},
}
b.writeMessage(server, conn)
}
func (b *Broker) writeFileOffset(offset int64, message *messages.Client, conn *websocket.Conn) {
msgType := messages.Server_FileOffset
server := &messages.Server{
MessageType: &msgType,
MessageId: proto.Uint64(message.GetMessageId()),
FileOffset: &messages.FileOffset{
Offset: proto.Int64(offset),
},
}
b.writeMessage(server, conn)
}
func (b *Broker) writeError(errStr string, message *messages.Client, conn *websocket.Conn) {
msgType := messages.Server_Error
server := &messages.Server{
MessageType: &msgType,
MessageId: proto.Uint64(message.GetMessageId()),
Error: &messages.Error{
Message: proto.String(errStr),
},
}
b.writeMessage(server, conn)
}
func (m *mockServer) readMessage(conn *websocket.Conn) *messages.Client {
_, data, err := conn.ReadMessage()
if err != nil {
return nil
}
var msg messages.Client
err = msg.Unmarshal(data)
Expect(err).ToNot(HaveOccurred())
return &msg
}
func (b *Broker) writeFileLocation(fetchFileErr *FetchFileError, message *messages.Client, conn *websocket.Conn) {
var uri string
if fetchFileErr != nil {
uri = fetchFileErr.Uri
}
msgType := messages.Server_FileLocation
server := &messages.Server{
MessageType: &msgType,
MessageId: proto.Uint64(message.GetMessageId()),
FileLocation: &messages.FileLocation{
Local: proto.Bool(uri == ""),
Uri: proto.String(uri),
},
}
b.writeMessage(server, conn)
}
server.Close()
})
It("returns an error", func(done Done) {
defer close(done)
mockServer.serverCh <- buildError(99, "some-error")
go func() {
defer GinkgoRecover()
err := connection.FetchFile(9, "some-file")
Expect(err).ToNot(BeNil())
Expect(err.Error()).To(Equal("some-error"))
}()
var clientMsg *messages.Client
Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
Expect(clientMsg.FetchFile.GetFileId()).To(BeEquivalentTo(9))
})
It("returns the file ID", func(done Done) {
defer close(done)
mockServer.serverCh <- buildFileLocation(99)
go func() {
defer GinkgoRecover()
err := connection.FetchFile(9, "some-file")
Expect(err).To(BeNil())
}
var wg sync.WaitGroup
defer wg.Wait()
wg.Add(len(mockServers))
for i := range mockServers {
go func(iter int) {
defer GinkgoRecover()
defer wg.Done()
_, err := client.FetchFile(fmt.Sprintf("some-file-%d", iter))
Expect(err).ToNot(HaveOccurred())
}(i)
}
for i, server := range mockServers {
var msg *messages.Client
Eventually(server.clientCh).Should(Receive(&msg))
Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))
Expect(msg.GetFetchFile().GetName()).To(Equal(fmt.Sprintf("some-file-%d", i)))
}
}, 5)
It("re-fetches a file if it is redirected", func(done Done) {
defer close(done)
mockServers[0].serverCh <- buildRemoteFileLocation(1, servers[1].URL)
mockServers[1].serverCh <- buildFileLocation(1)
_, err := client.FetchFile("some-file-1")
Expect(err).ToNot(HaveOccurred())
var msg *messages.Client