How to unit test a net.Conn function that modifies the message sent?

First, let me tell that I've seen other questions that are similar to this one, but I don't think any question really answers it in enough detail (How does one test net.Conn in unit tests in Golang? and How does one test net.Conn in unit tests in Golang?).

What I want to test is function that in some way responds to a TCP request.

In the simplest scenario:

func HandleMessage(c net.Conn) {
  s, err := bufio.NewReader(c).ReadString('\n')
  s = strings.Trim(s, "\n")

  c.Write([]byte(s + " whatever"))
}

How would I unit test such a function, ideally using net.Pipe() to not open actual connections. I've been trying things like this:

func TestHandleMessage(t *testing.T) {
  server, client := net.Pipe()

  go func(){
    server.Write([]byte("test"))
    server.Close()
  }()

  s, err := bufio.NewReader(c).ReadString('\n')

  assert.Nil(t, err)
  assert.Equal(t, "test whatever", s) 
}

However I can't seem to understand where to put the HandleMessage(client) (or HandleMessage(server)? in the test to actually get the response I want. I can get it to the point that it either blocks and won't finish at all, or that it will return the exact same string that I wrote in the server.

Could someone help me out and point out where I'm making a mistake? Or maybe point in a correct direction when it comes to testing TCP functions.

Solution

The net.Pipe docs say:

Pipe creates a synchronous, in-memory, full duplex network connection; both ends implement the Conn interface. Reads on one end are matched with writes on the other, copying data directly between the two; there is no internal buffering.

So the labels you are attaching to the net.Conn's (server and client) are arbitrary. If you find it simpler to understand feel free to use something line handleMessageConn, sendMessageConn := net.Pipe().

The below basically fills out the example given in the answer you mentioned.

func TestHandleMessage(t *testing.T) {
    server, client := net.Pipe()

    // Set deadline so test can detect if HandleMessage does not return
    client.SetDeadline(time.Now().Add(time.Second))

    // Configure a go routine to act as the server
    go func() {
        HandleMessage(server)
        server.Close()
    }()


    _, err := client.Write([]byte("test\n"))
    if err != nil {
        t.Fatalf("failed to write: %s", err)
    }

    // As the go routine closes the connection ReadAll is a simple way to get the response
    in, err := io.ReadAll(client)
    if err != nil {
        t.Fatalf("failed to read: %s", err)
    }

    // Using an Assert here will also work (if using a library that provides that functionality)
    if string(in) != "test whatever" {
        t.Fatalf("expected `test` got `%s`", in)
    }

    client.Close()
}

You could turn this around and put the Write/Read in the go routine but I believe the above approach is easier to understand and simplifies avoiding a limitation of the testing package:

A test ends when its Test function returns or calls any of the methods FailNow, Fatal, Fatalf, SkipNow, Skip, or Skipf. Those methods, as well as the Parallel method, must be called only from the goroutine running the Test function.

Note: If you don't need a net.Conn (as is the case in this simple example) consider using HandleMessage(c io.ReadWriter) (this provides users with more flexibility and also simplifies testing).