From 34c6a362e4af3dd08ecc6e5e36c837afa3ad82b5 Mon Sep 17 00:00:00 2001
From: Jaume Marhuenda <jaumemarhuenda@gmail.com>
Date: Tue, 16 Oct 2018 17:37:37 -0700
Subject: [PATCH] use ConnectTimeout for options message

---
 conn.go      | 4 +++-
 conn_test.go | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/conn.go b/conn.go
index 0bec55ccf..3414f5cd7 100644
--- a/conn.go
+++ b/conn.go
@@ -199,7 +199,6 @@ func (s *Session) dial(host *HostInfo, cfg *ConnConfig, errorHandler ConnErrorHa
 		r:             bufio.NewReader(conn),
 		cfg:           cfg,
 		calls:         make(map[int]*callReq),
-		timeout:       cfg.Timeout,
 		version:       uint8(cfg.ProtoVersion),
 		addr:          conn.RemoteAddr().String(),
 		errorHandler:  errorHandler,
@@ -232,11 +231,14 @@ func (s *Session) dial(host *HostInfo, cfg *ConnConfig, errorHandler ConnErrorHa
 		conn:        c,
 	}
 
+	c.timeout = cfg.ConnectTimeout
 	if err := startup.setupConn(ctx); err != nil {
 		c.close()
 		return nil, err
 	}
 
+	c.timeout = cfg.Timeout
+
 	// dont coalesce startup frames
 	if s.cfg.WriteCoalesceWaitTime > 0 {
 		c.w = newWriteCoalescer(c.w, s.cfg.WriteCoalesceWaitTime, c.quit)
diff --git a/conn_test.go b/conn_test.go
index 1a2adcbf2..afad0b769 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -658,7 +658,7 @@ func TestQueryTimeout(t *testing.T) {
 		if err != ErrTimeoutNoResponse {
 			t.Fatalf("expected to get %v for timeout got %v", ErrTimeoutNoResponse, err)
 		}
-	case <-time.After(10*time.Millisecond + db.cfg.Timeout):
+	case <-time.After(40*time.Millisecond + db.cfg.Timeout):
 		// ensure that the query goroutines have been scheduled
 		t.Fatalf("query did not timeout after %v", db.cfg.Timeout)
 	}