From a35f4ef32eff579f85ea95e7fdc5ae5f9b508c9c Mon Sep 17 00:00:00 2001
From: Amir Raminfar <findamir@gmail.com>
Date: Thu, 29 Nov 2018 19:36:49 -0800
Subject: [PATCH] Adds more tests

---
 __snapshots__/dozzle.snapshot | 11 +++++-
 main.go                       |  2 +-
 main_test.go                  | 63 +++++++++++++++++++++++++++++++----
 3 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/__snapshots__/dozzle.snapshot b/__snapshots__/dozzle.snapshot
index e598ed9f..4da5465d 100644
--- a/__snapshots__/dozzle.snapshot
+++ b/__snapshots__/dozzle.snapshot
@@ -3,4 +3,13 @@ HTTP/1.1 200 OK
 Connection: close
 Content-Type: text/plain; charset=utf-8
 
-[{"id":"1234567890","names":null,"name":"test","image":"image","imageId":"image_id","command":"command","created":0,"state":"state","status":"status"}]
\ No newline at end of file
+[{"id":"1234567890","names":null,"name":"test","image":"image","imageId":"image_id","command":"command","created":0,"state":"state","status":"status"}]
+
+/* snapshot: /api/logs/stream */
+HTTP/1.1 200 OK
+Connection: close
+Cache-Control: no-cache
+Connection: keep-alive
+Content-Type: text/event-stream
+
+data: INFO Testing logs...
\ No newline at end of file
diff --git a/main.go b/main.go
index e052e22c..dd212fc4 100644
--- a/main.go
+++ b/main.go
@@ -158,7 +158,7 @@ func (h *handler) streamLogs(w http.ResponseWriter, r *http.Request) {
         count := binary.BigEndian.Uint32(hdr[4:])
         _, err = io.CopyN(&buffer, reader, int64(count))
 
-        if err != nil {
+        if err != nil || err == io.EOF {
             log.Debugf("Error while reading from log stream: %v", err)
             break
         }
diff --git a/main_test.go b/main_test.go
index 16e785c1..62dac8f2 100644
--- a/main_test.go
+++ b/main_test.go
@@ -1,14 +1,20 @@
 package main
 
 import (
-    "github.com/amir20/dozzle/docker"
-    "github.com/beme/abide"
-    "github.com/stretchr/testify/mock"
-    "github.com/stretchr/testify/require"
+    "bytes"
+    "context"
+    "encoding/binary"
+    "io"
+    "io/ioutil"
     "net/http"
     "net/http/httptest"
     "os"
     "testing"
+
+    "github.com/amir20/dozzle/docker"
+    "github.com/beme/abide"
+    "github.com/stretchr/testify/mock"
+    "github.com/stretchr/testify/require"
 )
 
 type MockedClient struct {
@@ -18,12 +24,24 @@ type MockedClient struct {
 
 func (m *MockedClient) ListContainers() ([]docker.Container, error) {
     args := m.Called()
-    containers, _ := args.Get(0).([]docker.Container)
+    containers, ok := args.Get(0).([]docker.Container)
+    if !ok {
+        panic("containers is not of type []docker.Container")
+    }
     return containers, args.Error(1)
 }
 
-func Test_listContainers(t *testing.T) {
-    req, err := http.NewRequest("GET", "/health-check", nil)
+func (m *MockedClient) ContainerLogs(ctx context.Context, id string) (io.ReadCloser, error) {
+    args := m.Called(ctx, id)
+    reader, ok := args.Get(0).(io.ReadCloser)
+    if !ok {
+        panic("reader is not of type io.ReadCloser")
+    }
+    return reader, args.Error(1)
+}
+
+func Test_handler_listContainers(t *testing.T) {
+    req, err := http.NewRequest("GET", "/api/containers.json", nil)
     require.NoError(t, err, "NewRequest should not return an error.")
 
     rr := httptest.NewRecorder()
@@ -49,6 +67,37 @@ func Test_listContainers(t *testing.T) {
 
     handler.ServeHTTP(rr, req)
     abide.AssertHTTPResponse(t, "/api/containers.json", rr.Result())
+    mockedClient.AssertExpectations(t)
+}
+
+func Test_handler_streamLogs(t *testing.T) {
+    id := "123456"
+    req, err := http.NewRequest("GET", "/api/logs/stream", nil)
+    q := req.URL.Query()
+    q.Add("id", "123456")
+    req.URL.RawQuery = q.Encode()
+    require.NoError(t, err, "NewRequest should not return an error.")
+
+    rr := httptest.NewRecorder()
+
+    mockedClient := new(MockedClient)
+    log := "INFO Testing logs..."
+    b := make([]byte, 8)
+
+    binary.BigEndian.PutUint32(b[4:], uint32(len(log)))
+    b = append(b, []byte(log)...)
+
+    var reader io.ReadCloser
+    reader = ioutil.NopCloser(bytes.NewReader(b))
+    mockedClient.On("ContainerLogs", mock.Anything, id).Return(reader, nil)
+
+    h := handler{client: mockedClient}
+
+    handler := http.HandlerFunc(h.streamLogs)
+
+    handler.ServeHTTP(rr, req)
+    abide.AssertHTTPResponse(t, "/api/logs/stream", rr.Result())
+    mockedClient.AssertExpectations(t)
 }
 
 func TestMain(m *testing.M) {