From e53d5a02abd638be525d8898b4bd45d67b5d5a92 Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <yohann.danello@gmail.com>
Date: Wed, 14 Oct 2020 20:01:07 +0200
Subject: [PATCH] More coverage on telnet package

---
 stream/telnet/telnet_test.go | 46 ++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/stream/telnet/telnet_test.go b/stream/telnet/telnet_test.go
index 72c4116..cce5b69 100644
--- a/stream/telnet/telnet_test.go
+++ b/stream/telnet/telnet_test.go
@@ -6,10 +6,18 @@ import (
 	"math/rand"
 	"net"
 	"testing"
+	"time"
 )
 
 // TestTelnetOutput creates a TCP client that connects to the server and get one image.
 func TestTelnetOutput(t *testing.T) {
+	// Try to start Telnet server while it is disabled
+	Serve(&Options{Enabled: false})
+	StartASCIIArtStream("demo", ioutil.NopCloser(bytes.NewReader([]byte{})))
+	if GetNumberConnectedSessions("demo") != 0 {
+		t.Fatalf("Mysteriously found %d connected clients", GetNumberConnectedSessions("demo"))
+	}
+
 	// Enable and start Telnet server
 	Serve(&Options{
 		Enabled:       true,
@@ -32,6 +40,17 @@ func TestTelnetOutput(t *testing.T) {
 		t.Fatalf("Error while connecting to the TCP server: %s", err)
 	}
 
+	// Say goodbye
+	_, err = client.Write([]byte("exit"))
+	if err != nil {
+		t.Fatalf("Error while closing TCP connection: %s", err)
+	}
+
+	client, err = net.Dial("tcp", Cfg.ListenAddress)
+	if err != nil {
+		t.Fatalf("Error while connecting to the TCP server: %s", err)
+	}
+
 	// Create a sufficient large buffer
 	buff := make([]byte, 3*len(sampleImage))
 
@@ -46,6 +65,22 @@ func TestTelnetOutput(t *testing.T) {
 		t.Fatalf("Read %d bytes from TCP, expected %d, read: %s", n, len("[GHOSTREAM]\nEnter stream ID: "), buff[:n])
 	}
 
+	// Send wrong stream ID
+	_, err = client.Write([]byte("toto"))
+	if err != nil {
+		t.Fatalf("Error while writing from TCP: %s", err)
+	}
+	n, err = client.Read(buff)
+	if err != nil {
+		t.Fatalf("Error while reading from TCP: %s", err)
+	}
+	if n == len("Unknown stream ID.\n") {
+		_, err = client.Read(buff)
+		if err != nil {
+			t.Fatalf("Error while reading from TCP: %s", err)
+		}
+	}
+
 	// Send stream ID
 	_, err = client.Write([]byte("demo"))
 	if err != nil {
@@ -65,4 +100,15 @@ func TestTelnetOutput(t *testing.T) {
 	if GetNumberConnectedSessions("demo") != 1 {
 		t.Fatalf("Expected one telnet client only, found %d", GetNumberConnectedSessions("demo"))
 	}
+
+	// Close connection, ensure that the counter got decremented
+	err = client.Close()
+	if err != nil {
+		t.Fatalf("Error while closing telnet connection: %s", err)
+	}
+	// Wait for timeout
+	time.Sleep(time.Second)
+	if GetNumberConnectedSessions("demo") != 0 {
+		t.Fatalf("Expected no telnet client, found %d", GetNumberConnectedSessions("demo"))
+	}
 }