go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/milo/ui/src/proto_utils/batched_test_variant_branches_client.test.ts (about)

     1  // Copyright 2024 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  /* eslint-disable new-cap */
    16  
    17  import {
    18    BatchGetTestVariantBranchRequest,
    19    BatchGetTestVariantBranchResponse,
    20    TestVariantBranch,
    21    TestVariantBranchesClientImpl,
    22  } from '@/proto/go.chromium.org/luci/analysis/proto/v1/test_variant_branches.pb';
    23  
    24  import { BatchedTestVariantBranchesClientImpl } from './batched_test_variant_branches_client';
    25  
    26  describe('BatchedTestVariantBranchesClientImpl', () => {
    27    let batchGetSpy: jest.SpiedFunction<
    28      BatchedTestVariantBranchesClientImpl['BatchGet']
    29    >;
    30  
    31    beforeEach(() => {
    32      jest.useFakeTimers();
    33      batchGetSpy = jest
    34        .spyOn(TestVariantBranchesClientImpl.prototype, 'BatchGet')
    35        .mockImplementation(async (batchReq) => {
    36          return BatchGetTestVariantBranchResponse.fromPartial({
    37            testVariantBranches: Object.freeze(
    38              batchReq.names.map((name) =>
    39                TestVariantBranch.fromPartial({
    40                  name,
    41                }),
    42              ),
    43            ),
    44          });
    45        });
    46    });
    47  
    48    afterEach(() => {
    49      jest.useRealTimers();
    50      batchGetSpy.mockReset();
    51    });
    52  
    53    it('can batch eligible requests together', async () => {
    54      const client = new BatchedTestVariantBranchesClientImpl({
    55        request: jest.fn(),
    56      });
    57      const CALL_COUNT = 3;
    58      const BATCH_SIZE = 30;
    59  
    60      const calls = Array(CALL_COUNT)
    61        .fill(0)
    62        .map((_, i) =>
    63          client.BatchGet(
    64            BatchGetTestVariantBranchRequest.fromPartial({
    65              names: Object.freeze(
    66                Array(BATCH_SIZE)
    67                  .fill(0)
    68                  .map((_, j) => `call${i}-name${j}`),
    69              ),
    70            }),
    71          ),
    72        );
    73  
    74      await jest.advanceTimersToNextTimerAsync();
    75      // The batch function should be called only once.
    76      expect(batchGetSpy).toHaveBeenCalledTimes(1);
    77      expect(batchGetSpy).toHaveBeenCalledWith(
    78        BatchGetTestVariantBranchRequest.fromPartial({
    79          names: Object.freeze(
    80            Array(CALL_COUNT)
    81              .fill(0)
    82              .flatMap((_, i) =>
    83                Array(BATCH_SIZE)
    84                  .fill(0)
    85                  .map((_, j) => `call${i}-name${j}`),
    86              ),
    87          ),
    88        }),
    89      );
    90  
    91      // The responses should be just like regular calls.
    92      expect(await Promise.all(calls)).toEqual(
    93        Array(CALL_COUNT)
    94          .fill(0)
    95          .map((_, i) =>
    96            BatchGetTestVariantBranchResponse.fromPartial({
    97              testVariantBranches: Object.freeze(
    98                Array(BATCH_SIZE)
    99                  .fill(0)
   100                  .map((_, j) => ({ name: `call${i}-name${j}` })),
   101              ),
   102            }),
   103          ),
   104      );
   105    });
   106  
   107    it('can handle over batching', async () => {
   108      const client = new BatchedTestVariantBranchesClientImpl({
   109        request: jest.fn(),
   110      });
   111      const CALL_COUNT = 5;
   112      const BATCH_SIZE = 30;
   113  
   114      const calls = Array(CALL_COUNT)
   115        .fill(0)
   116        .map((_, i) =>
   117          client.BatchGet(
   118            BatchGetTestVariantBranchRequest.fromPartial({
   119              names: Object.freeze(
   120                Array(BATCH_SIZE)
   121                  .fill(0)
   122                  .map((_, j) => `call${i}-name${j}`),
   123              ),
   124            }),
   125          ),
   126        );
   127  
   128      await jest.advanceTimersToNextTimerAsync();
   129      // The batch function should be called more than once.
   130      expect(batchGetSpy).toHaveBeenCalledTimes(2);
   131  
   132      // The responses should be just like regular calls.
   133      expect(await Promise.all(calls)).toEqual(
   134        Array(CALL_COUNT)
   135          .fill(0)
   136          .map((_, i) =>
   137            BatchGetTestVariantBranchResponse.fromPartial({
   138              testVariantBranches: Object.freeze(
   139                Array(BATCH_SIZE)
   140                  .fill(0)
   141                  .map((_, j) => ({ name: `call${i}-name${j}` })),
   142              ),
   143            }),
   144          ),
   145      );
   146    });
   147  });