refactor(simplequicping): use step-by-step (#852)

See https://github.com/ooni/probe/issues/2159 and https://github.com/ooni/spec/pull/254

internal/measurexlite/conn_test.go

@@ -219,6 +219,174 @@ func TestWrapNetConn(t *testing.T) {
 	})
 }
 
+func TestWrapUDPLikeConn(t *testing.T) {
+	t.Run("WrapUDPLikeConn wraps the conn", func(t *testing.T) {
+		underlying := &mocks.UDPLikeConn{}
+		zeroTime := time.Now()
+		trace := NewTrace(0, zeroTime)
+		conn := trace.WrapUDPLikeConn(underlying)
+		ct := conn.(*udpLikeConnTrace)
+		if ct.UDPLikeConn != underlying {
+			t.Fatal("invalid underlying")
+		}
+		if ct.tx != trace {
+			t.Fatal("invalid trace")
+		}
+	})
+
+	t.Run("ReadFrom saves a trace", func(t *testing.T) {
+		underlying := &mocks.UDPLikeConn{
+			MockReadFrom: func(b []byte) (int, net.Addr, error) {
+				return len(b), &mocks.Addr{
+					MockString: func() string {
+						return "1.1.1.1:443"
+					},
+				}, nil
+			},
+		}
+		zeroTime := time.Now()
+		td := testingx.NewTimeDeterministic(zeroTime)
+		trace := NewTrace(0, zeroTime)
+		trace.TimeNowFn = td.Now // deterministic time counting
+		conn := trace.WrapUDPLikeConn(underlying)
+		const bufsiz = 128
+		buffer := make([]byte, bufsiz)
+		count, addr, err := conn.ReadFrom(buffer)
+		if count != bufsiz {
+			t.Fatal("invalid count")
+		}
+		if addr.String() != "1.1.1.1:443" {
+			t.Fatal("invalid address")
+		}
+		if err != nil {
+			t.Fatal("invalid err")
+		}
+		events := trace.NetworkEvents()
+		if len(events) != 1 {
+			t.Fatal("did not save network events")
+		}
+		expect := &model.ArchivalNetworkEvent{
+			Address:   "1.1.1.1:443",
+			Failure:   nil,
+			NumBytes:  bufsiz,
+			Operation: "read_from",
+			Proto:     "udp",
+			T:         1.0,
+			Tags:      []string{},
+		}
+		got := events[0]
+		if diff := cmp.Diff(expect, got); diff != "" {
+			t.Fatal(diff)
+		}
+	})
+
+	t.Run("ReadFrom discards the event when the buffer is full", func(t *testing.T) {
+		underlying := &mocks.UDPLikeConn{
+			MockReadFrom: func(b []byte) (int, net.Addr, error) {
+				return len(b), &mocks.Addr{
+					MockString: func() string {
+						return "1.1.1.1:443"
+					},
+				}, nil
+			},
+		}
+		zeroTime := time.Now()
+		trace := NewTrace(0, zeroTime)
+		trace.NetworkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
+		conn := trace.WrapUDPLikeConn(underlying)
+		const bufsiz = 128
+		buffer := make([]byte, bufsiz)
+		count, addr, err := conn.ReadFrom(buffer)
+		if count != bufsiz {
+			t.Fatal("invalid count")
+		}
+		if addr.String() != "1.1.1.1:443" {
+			t.Fatal("invalid address")
+		}
+		if err != nil {
+			t.Fatal("invalid err")
+		}
+		events := trace.NetworkEvents()
+		if len(events) != 0 {
+			t.Fatal("expected no network events")
+		}
+	})
+
+	t.Run("WriteTo saves a trace", func(t *testing.T) {
+		underlying := &mocks.UDPLikeConn{
+			MockWriteTo: func(b []byte, addr net.Addr) (int, error) {
+				return len(b), nil
+			},
+		}
+		zeroTime := time.Now()
+		td := testingx.NewTimeDeterministic(zeroTime)
+		trace := NewTrace(0, zeroTime)
+		trace.TimeNowFn = td.Now // deterministic time tracking
+		conn := trace.WrapUDPLikeConn(underlying)
+		const bufsiz = 128
+		buffer := make([]byte, bufsiz)
+		addr := &mocks.Addr{
+			MockString: func() string {
+				return "1.1.1.1:443"
+			},
+		}
+		count, err := conn.WriteTo(buffer, addr)
+		if count != bufsiz {
+			t.Fatal("invalid count")
+		}
+		if err != nil {
+			t.Fatal("invalid err")
+		}
+		events := trace.NetworkEvents()
+		if len(events) != 1 {
+			t.Fatal("did not save network events")
+		}
+		expect := &model.ArchivalNetworkEvent{
+			Address:   "1.1.1.1:443",
+			Failure:   nil,
+			NumBytes:  bufsiz,
+			Operation: "write_to",
+			Proto:     "udp",
+			T:         1.0,
+			Tags:      []string{},
+		}
+		got := events[0]
+		if diff := cmp.Diff(expect, got); diff != "" {
+			t.Fatal(diff)
+		}
+	})
+
+	t.Run("Write discards the event when the buffer is full", func(t *testing.T) {
+		underlying := &mocks.UDPLikeConn{
+			MockWriteTo: func(b []byte, addr net.Addr) (int, error) {
+				return len(b), nil
+			},
+		}
+		zeroTime := time.Now()
+		trace := NewTrace(0, zeroTime)
+		trace.NetworkEvent = make(chan *model.ArchivalNetworkEvent) // no buffer
+		conn := trace.WrapUDPLikeConn(underlying)
+		const bufsiz = 128
+		buffer := make([]byte, bufsiz)
+		addr := &mocks.Addr{
+			MockString: func() string {
+				return "1.1.1.1:443"
+			},
+		}
+		count, err := conn.WriteTo(buffer, addr)
+		if count != bufsiz {
+			t.Fatal("invalid count")
+		}
+		if err != nil {
+			t.Fatal("invalid err")
+		}
+		events := trace.NetworkEvents()
+		if len(events) != 0 {
+			t.Fatal("expected no network events")
+		}
+	})
+}
+
 func TestNewAnnotationArchivalNetworkEvent(t *testing.T) {
 	var (
 		index     int64 = 3