From f62470c0910efb1bce263aacb2e3b65c5837090a Mon Sep 17 00:00:00 2001
From: Alex Bramley <a.bramley@gmail.com>
Date: Fri, 11 Nov 2011 11:17:18 +0000
Subject: [PATCH] Test rateLimit(). Move call to time.After to write() for ease
 of testing.

Complete test coverage! Well, for things that matter. I think.
---
 client/connection.go      | 17 ++++++++++-------
 client/connection_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/client/connection.go b/client/connection.go
index 5b75c23..26dcf79 100644
--- a/client/connection.go
+++ b/client/connection.go
@@ -253,7 +253,12 @@ func (conn *Conn) runLoop() {
 // using Hybrid's algorithm to rate limit if conn.Flood is false.
 func (conn *Conn) write(line string) {
 	if !conn.Flood {
-		conn.rateLimit(int64(len(line)))
+		if t := conn.rateLimit(int64(len(line))); t != 0 {
+			// sleep for the current line's time value before sending it
+			conn.l.Debug("irc.rateLimit(): Flood! Sleeping for %.2f secs.",
+				float64(t)/float64(second))
+			<-time.After(t)
+		}
 	}
 
 	if _, err := conn.io.WriteString(line + "\r\n"); err != nil {
@@ -270,7 +275,7 @@ func (conn *Conn) write(line string) {
 }
 
 // Implement Hybrid's flood control algorithm to rate-limit outgoing lines.
-func (conn *Conn) rateLimit(chars int64) {
+func (conn *Conn) rateLimit(chars int64) int64 {
 	// Hybrid's algorithm allows for 2 seconds per line and an additional
 	// 1/120 of a second per character on that line.
 	linetime := 2*second + chars*second/120
@@ -282,12 +287,10 @@ func (conn *Conn) rateLimit(chars int64) {
 	conn.lastsent = time.Nanoseconds()
 	// If we've sent more than 10 second's worth of lines according to the
 	// calculation above, then we're at risk of "Excess Flood".
-	if conn.badness > 10*second && !conn.Flood {
-		// so sleep for the current line's time value before sending it
-		conn.l.Debug("irc.rateLimit(): Flood! Sleeping for %.2f secs.",
-			float64(linetime)/float64(second))
-		<-time.After(linetime)
+	if conn.badness > 10*second {
+		return linetime
 	}
+	return 0
 }
 
 func (conn *Conn) shutdown() {
diff --git a/client/connection_test.go b/client/connection_test.go
index e6d593f..479d9dd 100644
--- a/client/connection_test.go
+++ b/client/connection_test.go
@@ -373,3 +373,29 @@ func TestWrite(t *testing.T) {
 	s.log.EXPECT().Error("irc.send(): %s", "invalid argument")
 	c.write("she can't pass unit tests")
 }
+
+func TestRateLimit(t *testing.T) {
+	c, s := setUp(t)
+	defer s.tearDown()
+
+	if c.badness != 0 || c.lastsent != 0 {
+		t.Errorf("Bad initial values for rate limit variables.")
+	}
+
+	// badness will still be 0 because lastsent was 0 before rateLimit.
+	if l := c.rateLimit(60); l != 0 || c.badness != 0 || c.lastsent == 0 {
+		t.Errorf("Rate limit variables not updated correctly after rateLimit.")
+	}
+	// So, time at the nanosecond resolution is a bit of a bitch. Choosing 60
+	// characters as the line length means we should be increasing badness by
+	// 2.5 seconds minus the delta between the two ratelimit calls. This should
+	// be minimal but it's guaranteed that it won't be zero. Use 1us as a fuzz.
+	// This seems to be the minimum timer resolution, on my laptop at least...
+	if l := c.rateLimit(60); l != 0 || c.badness - int64(25*1e8) > 1e3 {
+		t.Errorf("Rate limit calculating badness incorrectly.")
+	}
+	// At this point, we can tip over the badness scale, with a bit of help.	
+	if l := c.rateLimit(360); l == 80*1e8 || c.badness - int64(105*1e8) > 1e3 {
+		t.Errorf("Rate limit failed to return correct limiting values.")
+	}
+}