Add more documentation and increase test coverage

broker.go

@@ -74,18 +74,21 @@ type AccessToken struct {
 	Token string
 	// Extensions is a optional map of arbitrary key-value pairs that can be
 	// sent with the SASL/OAUTHBEARER initial client response. These values are
-	// ignored by the SASL server if they are unexpected.
+	// ignored by the SASL server if they are unexpected. This feature is only
+	// supported by Kafka >= 2.1.0.
 	Extensions map[string]string
 }
 
 // AccessTokenProvider is the interface that encapsulates how implementors
 // can generate access tokens for Kafka broker authentication.
 type AccessTokenProvider interface {
-	// Token returns an access token. Because this method may be called multiple
-	// times, each invocation returns a new, unexpired token. This method should
-	// not block indefinitely. A timeout error should be returned after a short
-	// period of inactivity so that the broker connection logic can log
-	// debugging information and retry.
+	// Token returns an access token. The implementation should ensure token
+	// reuse so that multiple calls at connect time do not create multiple
+	// tokens. The implementation should also periodically refresh the token in
+	// order to guarantee that each call returns an unexpired token.  This
+	// method should not block indefinitely--a timeout error should be returned
+	// after a short period of inactivity so that the broker connection logic
+	// can log debugging information and retry.
 	Token() (*AccessToken, error)
 }
 

broker_test.go

@@ -177,7 +177,8 @@ func TestSASLOAuthBearer(t *testing.T) {
 		// mockBroker mocks underlying network logic and broker responses
 		mockBroker := NewMockBroker(t, 0)
 
-		mockSASLAuthResponse := NewMockSaslAuthenticateResponse(t)
+		mockSASLAuthResponse := NewMockSaslAuthenticateResponse(t).
+			SetAuthBytes([]byte(`response_payload`))
 
 		if test.mockAuthErr != ErrNoError {
 			mockSASLAuthResponse = mockSASLAuthResponse.SetError(test.mockAuthErr)

mockresponses.go

@@ -707,8 +707,9 @@ func (mr *MockListAclsResponse) For(reqBody versionedDecoder) encoder {
 }
 
 type MockSaslAuthenticateResponse struct {
-	t      TestReporter
-	kerror KError
+	t             TestReporter
+	kerror        KError
+	saslAuthBytes []byte
 }
 
 func NewMockSaslAuthenticateResponse(t TestReporter) *MockSaslAuthenticateResponse {
@@ -718,6 +719,7 @@ func NewMockSaslAuthenticateResponse(t TestReporter) *MockSaslAuthenticateRespon
 func (msar *MockSaslAuthenticateResponse) For(reqBody versionedDecoder) encoder {
 	res := &SaslAuthenticateResponse{}
 	res.Err = msar.kerror
+	res.SaslAuthBytes = msar.saslAuthBytes
 	return res
 }
 
@@ -726,6 +728,11 @@ func (msar *MockSaslAuthenticateResponse) SetError(kerror KError) *MockSaslAuthe
 	return msar
 }
 
+func (msar *MockSaslAuthenticateResponse) SetAuthBytes(saslAuthBytes []byte) *MockSaslAuthenticateResponse {
+	msar.saslAuthBytes = saslAuthBytes
+	return msar
+}
+
 type MockDeleteAclsResponse struct {
 	t TestReporter
 }