go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/appengine/tq/internals.go (about)

     1  // Copyright 2017 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package tq
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/golang/protobuf/proto"
    21  
    22  	"go.chromium.org/luci/common/data/stringset"
    23  	"go.chromium.org/luci/gae/service/taskqueue"
    24  )
    25  
    26  // Internals is used by tqtesting package and must not be used directly.
    27  //
    28  // For that reason it returns opaque interface type, to curb the curiosity.
    29  //
    30  // We do this to avoid linking testing implementation into production binaries.
    31  // We make testing live in a different package and use this secret back door API
    32  // to talk to Dispatcher.
    33  func (d *Dispatcher) Internals() any {
    34  	return internalsImpl{d}
    35  }
    36  
    37  // internalsImpl secretly conforms to tqtesting.dispatcherInternals interface.
    38  type internalsImpl struct {
    39  	*Dispatcher
    40  }
    41  
    42  func (d internalsImpl) GetBaseURL() string {
    43  	return d.baseURL()
    44  }
    45  
    46  func (d internalsImpl) GetAllQueues() []string {
    47  	qs := stringset.New(0)
    48  	d.mu.RLock()
    49  	for _, h := range d.handlers {
    50  		qs.Add(h.queue)
    51  	}
    52  	d.mu.RUnlock()
    53  	return qs.ToSlice()
    54  }
    55  
    56  func (d internalsImpl) GetPayload(blob []byte) (proto.Message, error) {
    57  	payload, err := deserializePayload(blob)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	if _, err := d.handler(payload); err != nil {
    62  		return nil, err
    63  	}
    64  	return payload, nil
    65  }
    66  
    67  func (d internalsImpl) GetHandler(payload proto.Message) (cb Handler, q string, err error) {
    68  	h, err := d.handler(payload)
    69  	if err != nil {
    70  		return nil, "", err
    71  	}
    72  	return h.cb, h.queue, nil
    73  }
    74  
    75  func (d internalsImpl) WithRequestHeaders(ctx context.Context, hdr *taskqueue.RequestHeaders) context.Context {
    76  	return withRequestHeaders(ctx, hdr)
    77  }