Function testProxyRetryWithConfig

proxy/proxy_retries_test.go:373–472

(t *testing.T, query *frame.Frame, response message.Error, cfg *proxyTestConfig, testMessage string)

}

func testProxyRetryWithConfig(t *testing.T, query *frame.Frame, response message.Error, cfg *proxyTestConfig, testMessage string) (numNodesTried, retryCount int, responseError error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var mu sync.Mutex
	tried := make(map[string]int)
	prepared := make(map[[16]byte]string)

	if cfg == nil {
		cfg = &proxyTestConfig{}
	}

	cfg.handlers = proxycore.MockRequestHandlers{
		primitive.OpCodeQuery: func(cl *proxycore.MockClient, frm *frame.Frame) message.Message {
			if msg := cl.InterceptQuery(frm.Header, frm.Body.Message.(*message.Query)); msg != nil {
				return msg
			} else {
				mu.Lock()
				tried[cl.Local().IP]++
				mu.Unlock()
				return response
			}
		},
		primitive.OpCodeExecute: func(cl *proxycore.MockClient, frm *frame.Frame) message.Message {
			msg := frm.Body.Message.(*message.Execute)
			mu.Lock()
			defer mu.Unlock()
			var id [16]byte
			copy(id[:], msg.QueryId)
			if _, ok := prepared[id]; !ok {
				return &message.Unprepared{
					ErrorMessage: "query is not prepared",
					Id:           id[:],
				}
			} else {
				tried[cl.Local().IP]++
				return response
			}
		},
		primitive.OpCodeBatch: func(cl *proxycore.MockClient, frm *frame.Frame) message.Message {
			msg := frm.Body.Message.(*message.Batch)
			mu.Lock()
			defer mu.Unlock()
			for _, child := range msg.Children {
				id := child.Id
				if id != nil {
					var hash [16]byte
					copy(hash[:], id)
					if _, ok := prepared[hash]; !ok {
						return &message.Unprepared{
							ErrorMessage: "query is not prepared",
							Id:           id,
						}
					}
				}
			}
			tried[cl.Local().IP]++
			return response