Handling imprecise timing when testing cancellation of paged query - go

Question

I have an object I'm using to make paged SQL queries that allows for the queries to be run asynchronously:
type PagedQuery[T any] struct {
Results chan []*T
Errors chan error
Done chan error
Quit chan error
client *sql.DB
}
func NewPagedQuery[T any](client *sql.DB) *PagedQuery[T] {
return &PagedQuery[T]{
Results: make(chan []*T, 1),
Errors: make(chan error, 1),
Done: make(chan error, 1),
Quit: make(chan error, 1),
client: client,
}
}
func (paged *PagedQuery[T]) requestAsync(ctx context.Context, queries ...*Query) {
conn, err := client.Conn(ctx)
if err != nil {
paged.Errors <- err
return
}
defer func() {
conn.Close()
paged.Done <- nil
}()
for i, query := range queries {
select {
case <-ctx.Done():
return
case <-paged.Quit:
return
default:
}
rows, err := conn.QueryContext(ctx, query.String, query.Arguments...)
if err != nil {
paged.Errors <- err
return
}
data, err := sql.ReadRows[T](rows)
if err != nil {
paged.Errors <- err
return
}
paged.Results <- data
}
}
I'm trying to test this code, specifically the cancellation part. My test code looks like this:
svc, mock := createServiceMock("TEST_DATABASE", "TEST_SCHEMA")
mock.ExpectQuery(regexp.QuoteMeta("TEST QUERY")).
WithArgs(...).
WillReturnRows(mock.NewRows([]string{"t", "v", "o", "c", "h", "l", "vw", "n"}))
ctx, cancel := context.WithCancel(context.Background())
go svc.requestAsync(ctx, query1, query2, query3, query4)
time.Sleep(50 * time.Millisecond)
cancel()
results := make([]data, 0)
loop:
for {
select {
case <-query.Done:
break loop
case err := <-query.Errors:
Expect(err).ShouldNot(HaveOccurred())
case r := <-query.Results:
results = append(results, r...)
}
}
Expect(results).Should(BeEmpty())
Expect(mock.ExpectationsWereMet()).ShouldNot(HaveOccurred()) // fails here
The issue I'm running into is that this test fails occaisionally at the line indicated by my comment, because when cancel() is called, execution isn't guaranteed to be at the switch statement where I check for <-ctx.Done or <-Quit. Execution could be anywhere in the loop up to where I send the results to the Results channel. Except that doesn't make sense because execution should block until I receive data from the Results channel, which I don't do until after I call cancel(). Furthermore, I'm relying on the sqlmock package for SQL testing which doesn't allow for any sort of fuzzy checking where SQL queries are concerned. Why am I getting this failure and how can I fix it?

Answer 1

My issue was the result of my own lack of understanding around Go channels. I assumed that, by creating a chan([]*T, 1) meant that the channel would block when it was full (i.e. when it contained a single item) but that is not the case. Rather, the block occurs when I attempted to send to the channel when its buffer was full. So, by modifying Results like this:
func NewPagedQuery[T any](client *sql.DB) *PagedQuery[T] {
return &PagedQuery[T]{
Results: make(chan []*T), // Remove buffer here
Errors: make(chan error, 1),
Done: make(chan error, 1),
Quit: make(chan error, 1),
client: client,
}
}
I can ensure that the channel blocks until the data it contains is received. This one change fixed all the problems with testing.