Proper way to test a connection in go - unit-testing

I am covering project with tests and for that purpose I need dummy TCP Server, which could accept connection, write/read data to/from it, close it etc... I have found this question on stack overflow, covering mocking connection, but it doesn't cover what I actually need to test.
My idea relies on this article as starting point, but when I started implementing channel to let server write some data to newly opened connection, I got stuck with undebuggable deadlock in writing to channel.
What I want to achieve is to write some data to server's channel, say sendingQueue chan *[]byte, so later corresponding []byte will be sent to newly established connection.
During these little research I have tried debugging and printing out messages before/after sending data to channel and trying to send / read data from channel in different places of program.
What I found out:
My idea works if I add data directly in handleConnection with
go func() {
f := []byte("test.")
t.sendingQueue <- &f
}()
My idea doesn't work if I push data to channel from TestUtils_TestingTCPServer_WritesRequest in any form, either with func (t *TCPServer) Put(data *[]byte) (err error) or directly with:
go func(queue chan *[]byte, data *[]byte) {
queue <- data
}(t.sendingQueue, &payload)
It doesn't matter if channel is buffered or not.
So, obviously, there is something wrong either with the way I debug my code (I didn't dive into cli dlv, using just IDE debugger), or something that I completely miss about working with go channels, goroutines or net.Conn module.
For convenience public gist with full code is available. Note — there is // INIT part in the TestUtils_TestingTCPServer_WritesRequest which is required to run/debug single test. It should be commented out when running go test in the directory.
utils.go:
// NewServer creates a new Server using given protocol
// and addr.
func NewTestingTCPServer(protocol, addr string) (*TCPServer, error) {
switch strings.ToLower(protocol) {
case "tcp":
return &TCPServer{
addr: addr,
sendingQueue: make(chan *[]byte, 10),
}, nil
case "udp":
}
return nil, errors.New("invalid protocol given")
}
// TCPServer holds the structure of our TCP
// implementation.
type TCPServer struct {
addr string
server net.Listener
sendingQueue chan *[]byte
}
func (t *TCPServer) Run() (err error) {}
func (t *TCPServer) Close() (err error) {}
func (t *TCPServer) Put(data *[]byte) (err error) {}
func (t *TCPServer) handleConnection(conn net.Conn){
// <...>
// Putting data here successfully sends it via freshly established
// Connection:
// go func() {
// f := []byte("test.")
// t.sendingQueue <- &f
// }()
for {
fmt.Printf("Started for loop\n")
select {
case data := <-readerChannel:
fmt.Printf("Read written data\n")
writeBuffer.Write(*data)
writeBuffer.Flush()
case data := <-t.sendingQueue:
fmt.Printf("Read pushed data\n")
writeBuffer.Write(*data)
writeBuffer.Flush()
case <-ticker:
fmt.Printf("Tick\n")
return
}
fmt.Printf("Finished for loop\n")
}
}
utils_test.go
func TestUtils_TestingTCPServer_WritesRequest(t *testing.T) {
payload := []byte("hello world\n")
// <...> In gist here is placed INIT piece, which
// is required to debug single test
fmt.Printf("Putting payload into queue\n")
// This doesn't affect channel
err = utilTestingSrv.Put(&payload)
assert.Nil(t, err)
// This doesn't work either
//go func(queue chan *[]byte, data *[]byte) {
// queue <- data
//}(utilTestingSrv.sendingQueue, &payload)
conn, err := net.Dial("tcp", ":41123")
if !assert.Nil(t, err) {
t.Error("could not connect to server: ", err)
}
defer conn.Close()
out := make([]byte, 1024)
if _, err := conn.Read(out); assert.Nil(t, err) {
// Need to remove trailing byte 0xa from bytes array to make sure bytes array are equal.
if out[len(payload)] == 0xa {
out[len(payload)] = 0x0
}
assert.Equal(t, payload, bytes.Trim(out, "\x00"))
} else {
t.Error("could not read from connection")
}
}

After a help from a colleague and reading the article on how init works, I found a problem.
It was in init function, which was recreating extra server, due to using := assignment. I also updated code to make sure server runs before net.Dial and conn.Read.

Related

Unit Testing an SSH Client in Go

I wrote an SSH client in Go and I would like to write some tests. The problem is that I've never really written proper unit tests before, and most tutorials seem to focus on writing tests for a function that adds two numbers or some other toy problem. I've read about mocking, using interfaces, and other techniques, but I'm having trouble applying them. Also, my client is going to be used concurrently to allow fast configuration of multiple devices at a time. Not sure if that would change the way I write my tests or would add additional tests. Any help is appreciated.
Here is my code. Basically, a Device has 4 main functions: Connect, Send, Output/Err and Close for connecting to a device, sending it a set of configuration commands, capturing the output of the session, and closing the client, respectively.
package device
import (
"bufio"
"fmt"
"golang.org/x/crypto/ssh"
"io"
"net"
"time"
)
// A Device represents a remote network device.
type Device struct {
Host string // the device's hostname or IP address
client *ssh.Client // the client connection
session *ssh.Session // the connection to the remote shell
stdin io.WriteCloser // the remote shell's standard input
stdout io.Reader // the remote shell's standard output
stderr io.Reader // the remote shell's standard error
}
// Connect establishes an SSH connection to a device and sets up the session IO.
func (d *Device) Connect(user, password string) error {
// Create a client connection
client, err := ssh.Dial("tcp", net.JoinHostPort(d.Host, "22"), configureClient(user, password))
if err != nil {
return err
}
d.client = client
// Create a session
session, err := client.NewSession()
if err != nil {
return err
}
d.session = session
return nil
}
// configureClient sets up the client configuration for login
func configureClient(user, password string) *ssh.ClientConfig {
var sshConfig ssh.Config
sshConfig.SetDefaults()
sshConfig.Ciphers = append(sshConfig.Ciphers, "aes128-cbc", "aes256-cbc", "3des-cbc", "des-cbc", "aes192-cbc")
config := &ssh.ClientConfig{
Config: sshConfig,
User: user,
Auth: []ssh.AuthMethod{ssh.Password(password)},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
Timeout: time.Second * 5,
}
return config
}
// setupIO creates the pipes connected to the remote shell's standard input, output, and error
func (d *Device) setupIO() error {
// Setup standard input pipe
stdin, err := d.session.StdinPipe()
if err != nil {
return err
}
d.stdin = stdin
// Setup standard output pipe
stdout, err := d.session.StdoutPipe()
if err != nil {
return err
}
d.stdout = stdout
// Setup standard error pipe
stderr, err := d.session.StderrPipe()
if err != nil {
return err
}
d.stderr = stderr
return nil
}
// Send sends cmd(s) to the device's standard input. A device only accepts one call
// to Send, as it closes the session and its standard input pipe.
func (d *Device) Send(cmds ...string) error {
if d.session == nil {
return fmt.Errorf("device: session is closed")
}
defer d.session.Close()
// Start the shell
if err := d.startShell(); err != nil {
return err
}
// Send commands
for _, cmd := range cmds {
if _, err := d.stdin.Write([]byte(cmd + "\r")); err != nil {
return err
}
}
defer d.stdin.Close()
// Wait for the commands to exit
d.session.Wait()
return nil
}
// startShell requests a pseudo terminal (VT100) and starts the remote shell.
func (d *Device) startShell() error {
modes := ssh.TerminalModes{
ssh.ECHO: 0, // disable echoing
ssh.OCRNL: 0,
ssh.TTY_OP_ISPEED: 14400,
ssh.TTY_OP_OSPEED: 14400,
}
err := d.session.RequestPty("vt100", 0, 0, modes)
if err != nil {
return err
}
if err := d.session.Shell(); err != nil {
return err
}
return nil
}
// Output returns the remote device's standard output output.
func (d *Device) Output() ([]string, error) {
return readPipe(d.stdout)
}
// Err returns the remote device's standard error output.
func (d *Device) Err() ([]string, error) {
return readPipe(d.stdout)
}
// reapPipe reads an io.Reader line by line
func readPipe(r io.Reader) ([]string, error) {
var lines []string
scanner := bufio.NewScanner(r)
for scanner.Scan() {
lines = append(lines, scanner.Text())
}
if err := scanner.Err(); err != nil {
return nil, err
}
return lines, nil
}
// Close closes the client connection.
func (d *Device) Close() error {
return d.client.Close()
}
// String returns the string representation of a `Device`.
func (d *Device) String() string {
return fmt.Sprintf("%s", d.Host)
}
You make a good point about unit test tutorials nearly always being toy problems (why is it always Fibonacci?), when what we have is databases and http servers. The big realization that helped me is that you can only unit test things where you can control the input and output of the unit. configureClient or readPipe (give it a strings.Reader) would be good candidates. Start there.
Anything that leaves your program by talking directly to the disk, the network, stdout, etc, like the Connect method you would consider part of the external interface of your program. You don't unit test those. You integration test them.
Change Device to be an interface rather than a struct, and make a MockDevice that implements it. The real device is now maybe SSHDevice. You can unit test the rest of your program (which uses Device interface) by inserting a MockDevice, to isolate yourself from the network.
The SSHDevice will get tested in your integration tests. Start a real ssh server (maybe a test one you write in Go using crypto/ssh package, but any sshd would work). Start your program with an SSHDevice, make them talk to each other, and check outputs. You'll be using the os/exec package a lot. Integration tests are even more fun to write than unit tests!

Golang upload Http request FormFile to Amazon S3

I'm creating a micro service to handle some attachments uploads to Amazon S3, What I'm trying to achieve is accept a file and then store it directly to my Amazon S3 bucket, my current function :
func upload_handler(w http.ResponseWriter, r *http.Request) {
file, header, err := r.FormFile("attachment")
if err != nil {
fmt.Fprintln(w, err)
return
}
defer file.Close()
fileSize, err := file.Seek(0, 2) //2 = from end
if err != nil {
panic(err)
}
fmt.Println("File size : ", fileSize)
bytes := make([]byte, fileSize)
// read into buffer
buffer := bufio.NewReader(file)
_, err = buffer.Read(bytes)
auth := aws.Auth{
AccessKey: "XXXXXXXXXXX",
SecretKey: "SECRET_KEY_HERE",
}
client := s3.New(auth, aws.EUWest)
bucket := client.Bucket("attachments")
err = bucket.Put(header.Filename, bytes, header.Header.Get("Content-Type"), s3.ACL("public-read"))
if err != nil {
fmt.Println(err)
os.Exit(1)
}
}
The problem is that the files stored in S3 are all corrupted, After a small verification it seems that the file payload is not read as bytes
How to convert the file to bytes and store it correctly to S3 ?
Use ioutil.ReadAll:
bs, err := ioutil.ReadAll(file)
// ...
err = bucket.Put(
header.Filename,
bs,
header.Header.Get("Content-Type"),
s3.ACL("public-read"),
)
Read is a lower-level function which has subtle behavior:
Read reads data into p. It returns the number of bytes read into p. It calls Read at most once on the underlying Reader, hence n may be less than len(p). At EOF, the count will be zero and err will be io.EOF.
So what was probably happening was some subset of the file data was being written to S3 along with a bunch of 0s.
ioutil.ReadAll works by calling Read over and over again filling a dynamically expanding buffer until it reaches the end of the file. (so there's no need for the bufio.Reader either)
Also the Put function will have issues with large files (using ReadAll means the entire file must fit in memory) so you may want to use PutReader instead:
bucket.PutReader(
header.Filename,
file,
fileSize,
header.Header.Get("Content-Type"),
s3.ACL("public-read"),
)

How does one test net.Conn in unit tests in Golang?

I'm currently looking into creating some unit tests for net.Conn interface in Go, as well as other functions that build up on top of that functionality, and I'm wondering what is the best way to unit test that in Google Go? My code looks like:
conn, _:=net.Dial("tcp", "127.0.0.1:8080")
...
fmt.Fprintf(conn, "test")
...
buffer:=make([]byte, 100)
conn.Read(buffer)
Is the most efficient way of testing this code and the code that uses these functions to spin up a separate goroutine to act like the server, use net.http.httptest package, or something else?
You might be able to do what you need with net.Pipe which basically gives you both ends of a connection (think, after .Accept())
server, client := net.Pipe()
go func() {
// Do some stuff
server.Close()
}()
// Do some stuff
client.Close()
Although it will depend on the implementation details of your particular case, the general approach will be to start a server (in a separate goroutine, as you already hinted), and listen to the incoming connections.
For example, let's spin up a server and verify that the content we are reading from the connection is indeed the one we send over from the client:
func TestConn(t *testing.T) {
message := "Hi there!\n"
go func() {
conn, err := net.Dial("tcp", ":3000")
if err != nil {
t.Fatal(err)
}
defer conn.Close()
if _, err := fmt.Fprintf(conn, message); err != nil {
t.Fatal(err)
}
}()
l, err := net.Listen("tcp", ":3000")
if err != nil {
t.Fatal(err)
}
defer l.Close()
for {
conn, err := l.Accept()
if err != nil {
return
}
defer conn.Close()
buf, err := ioutil.ReadAll(conn)
if err != nil {
t.Fatal(err)
}
fmt.Println(string(buf[:]))
if msg := string(buf[:]); msg != message {
t.Fatalf("Unexpected message:\nGot:\t\t%s\nExpected:\t%s\n", msg, message)
}
return // Done
}
}
Note that here I'm not starting the server in the goroutine, as otherwise the test case is likely to be finished before the listener has run the test.
Another option is the counterfeiter package which lets you create mocks from interfaces and then you can stub out whatever calls you need. I have used it with great success to stub out net.Conn instances where I am testing out a protobuf client for Geode.
For example - https://github.com/gemfire/geode-go-client/blob/master/connector/protobuf_test.go

Ordering of senders in a Go channel

Consider the ping pong example from http://www.golang-book.com/10/index.htm#section2.
package main
import (
"fmt"
"time"
)
func pinger(c chan string) {
for i := 0; ; i++ {
c <- "ping"
}
}
func ponger(c chan string) {
for i := 0; ; i++ {
c <- "pong"
}
}
func printer(c chan string) {
for {
msg := <- c
fmt.Println(msg)
time.Sleep(time.Second * 1)
}
}
func main() {
var c chan string = make(chan string)
go pinger(c)
go ponger(c)
go printer(c)
var input string
fmt.Scanln(&input)
}
The authors write:
"The program will now take turns printing ping and pong."
However, for this to be true, Go must decide on an order in which senders can send into a channel? Otherwise, there would be no guarantee that a ping would be sent before a pong (i.e. you can't get two pings, or two pongs in a row). How does this work?
There is no synchronization between the ping and pong goroutines, therefore there is no guarantee that the responses will print in order.
If you force the goroutines to race with GOMAXPROCS>1, you get random output:
pong
ping
ping
pong
ping
pong
ping
pong
pong
This isn't even an example of a "ping-pong", since there's is no call and response.
There was a related question on the order of selection of messages entering a channel recently.
The answer is that the order is normally non-deterministic. This is intentional.

How do I read a UDP connection until a timeout is reached?

I need to read UDP traffic until a timeout is reached. I can do this by calling SetDeadline on the UDPConn and looping until I get an I/O timeout error, but this seems hack-ish (flow control based on error conditions). The following code snippet seems more correct, but does not terminate. In production, this would obviously be executed in a goroutine; it's written as a main function for simplicity's sake.
package main
import (
"fmt"
"time"
)
func main() {
for {
select {
case <-time.After(time.Second * 1):
fmt.Printf("Finished listening.\n")
return
default:
fmt.Printf("Listening...\n")
//read from UDPConn here
}
}
}
Why doesn't the given program terminate? Based on https://gobyexample.com/select, https://gobyexample.com/timeouts, and https://gobyexample.com/non-blocking-channel-operations, I would expect the above code to select the default case for one second, then take the first case and break out of the loop. How might I modify the above snippet to achieve the desired effect of looping and reading until a timeout occurs?
If you are not concerned about read blocking past n seconds, then loop until the deadline:
deadline := time.Now().Add(n * time.Second)
for time.Now().Before(deadline) {
fmt.Printf("Listening...\n")
//read from UDPConn here
}
fmt.Printf("Finished listening.\n")
If you do want to break out of a blocking read after n seconds, then set a deadline and read until there's an error:
conn.SetReadDeadline(time.Now().Add(n * time.Second)
for {
n, err := conn.Read(buf)
if err != nil {
if e, ok := err.(net.Error); !ok || !e.Timeout() {
// handle error, it's not a timeout
}
break
}
// do something with packet here
}
Using a deadline is not hacky. The standard library uses deadlines while reading UDP connections (see the dns client).
There are alternatives to using a deadline to break a blocking read: close the connection or send a dummy packet that the reader recognizes. These alternatives require starting another goroutine and are much more complicated than setting a deadline.
Simply assign the channel from time.After outside the for loop, otherwise you will just create a new timer each time you loop.
Example:
func main() {
ch := time.After(time.Second * 1)
L:
for {
select {
case <-ch:
fmt.Printf("Finished listening.\n")
break L // have to use a label or it will just break select
default:
fmt.Printf("Listening...\n")
//read from UDPConn here
}
}
}
Note that this doesn't work on the playground.