From 2fc42765180fd5341dad0a5136a51c035ae71ad9 Mon Sep 17 00:00:00 2001
From: Eygin Semen Leonidovich <seygin@ozon.ru>
Date: Wed, 9 Nov 2022 18:57:08 +0300
Subject: [PATCH] query tests added

---
 queries_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/queries_test.go b/queries_test.go
index 2b1125df..79dea168 100644
--- a/queries_test.go
+++ b/queries_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -1370,6 +1371,46 @@ func TestProcessQueryErrors(t *testing.T) {
 	}
 }
 
+type mockReadWriteCloser struct {
+	io.ReadWriteCloser
+}
+
+func (*mockReadWriteCloser) Read([]byte) (int, error) { return 0, errors.New("fake err") }
+
+func TestProcessQueryCancelConfirmationError(t *testing.T) {
+	tl := testLogger{t: t}
+	defer tl.StopLogging()
+	conn := internalConnection(t, &tl)
+	defer conn.Close()
+
+	stmt, err := conn.prepareContext(context.Background(), "select 1")
+	if err != nil {
+		t.Fatal("prepareContext expected to succeed, but it failed with", err)
+	}
+	err = stmt.sendQuery(context.Background(), []namedValue{})
+	if err != nil {
+		t.Fatal("sendQuery expected to succeed, but it failed with", err)
+	}
+	// mock real connection to imitate situation when you write but dont get response
+	conn.sess.buf.transport = &mockReadWriteCloser{ReadWriteCloser: conn.sess.buf.transport}
+	// canceling context to try to send attention request
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err = stmt.processQueryResponse(ctx)
+	if err == nil {
+		t.Error("processQueryResponse expected to fail but it succeeded")
+	}
+	// should not fail with ErrBadConn because query was successfully sent to server
+	if err != ErrorCancelConfirmation {
+		t.Error("processQueryResponse expected to fail with ErrorCancelConfirmation error but failed with other error: ", err)
+	}
+
+	if conn.connectionGood {
+		t.Fatal("Connection should be in a bad state")
+	}
+}
+
 func TestProcessQueryNextErrors(t *testing.T) {
 	tl := testLogger{t: t}
 	defer tl.StopLogging()