package transport

import (
	"net"
	"testing"
	"time"
)

type recordingConn struct {
	writeDeadlines []time.Time
}

func (c *recordingConn) Read(_ []byte) (int, error)        { return 0, nil }
func (c *recordingConn) Write(b []byte) (int, error)       { return len(b), nil }
func (c *recordingConn) Close() error                      { return nil }
func (c *recordingConn) LocalAddr() net.Addr               { return noopAddr("local") }
func (c *recordingConn) RemoteAddr() net.Addr              { return noopAddr("remote") }
func (c *recordingConn) SetDeadline(_ time.Time) error     { return nil }
func (c *recordingConn) SetReadDeadline(_ time.Time) error { return nil }
func (c *recordingConn) SetWriteDeadline(t time.Time) error {
	c.writeDeadlines = append(c.writeDeadlines, t)
	return nil
}

type noopAddr string

func (a noopAddr) Network() string { return string(a) }
func (a noopAddr) String() string  { return string(a) }

func TestWriteDeadlineConnSetsAndClearsDeadline(t *testing.T) {
	base := &recordingConn{}
	conn := &writeDeadlineConn{Conn: base, timeout: 25 * time.Millisecond}

	n, err := conn.Write([]byte("ping"))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	if n != 4 {
		t.Fatalf("Write bytes: got %d want 4", n)
	}
	if len(base.writeDeadlines) != 2 {
		t.Fatalf("write deadline calls: got %d want 2", len(base.writeDeadlines))
	}
	if base.writeDeadlines[0].IsZero() {
		t.Fatal("first write deadline must be non-zero")
	}
	if !base.writeDeadlines[1].IsZero() {
		t.Fatalf("second write deadline must clear the deadline, got %v", base.writeDeadlines[1])
	}
}