src/csharp/Grpc.HealthCheck.Tests/TestResponseStreamWriter.cs

   1 #region Copyright notice and license
   2 // Copyright 2015 gRPC authors.
   3 //
   4 // Licensed under the Apache License, Version 2.0 (the "License");
   5 // you may not use this file except in compliance with the License.
   6 // You may obtain a copy of the License at
   7 //
   8 //     http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15 #endregion
  16
  17 #if GRPC_SUPPORT_WATCH
  18 using System.Threading.Channels;
  19 using System.Threading.Tasks;
  20
  21 using Grpc.Core;
  22 using Grpc.Health.V1;
  23
  24 namespace Grpc.HealthCheck.Tests
  25 {
  26     internal class TestResponseStreamWriter : IServerStreamWriter<HealthCheckResponse>
  27     {
  28         private readonly Channel<HealthCheckResponse> _channel;
  29         private readonly TaskCompletionSource<object> _startTcs;
  30
  31         public TestResponseStreamWriter(int maxCapacity = 1, bool started = true)
  32         {
  33             _channel = System.Threading.Channels.Channel.CreateBounded<HealthCheckResponse>(new BoundedChannelOptions(maxCapacity) {
  34                 SingleReader = false,
  35                 SingleWriter = true,
  36                 FullMode = BoundedChannelFullMode.Wait
  37             });
  38             if (!started)
  39             {
  40                 _startTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
  41             }
  42         }
  43
  44         public ChannelReader<HealthCheckResponse> WrittenMessagesReader => _channel.Reader;
  45
  46         public WriteOptions WriteOptions { get; set; }
  47
  48         public async Task WriteAsync(HealthCheckResponse message)
  49         {
  50             if (_startTcs != null)
  51             {
  52                 await _startTcs.Task;
  53             }
  54
  55             await _channel.Writer.WriteAsync(message);
  56         }
  57
  58         public void Start()
  59         {
  60             if (_startTcs != null)
  61             {
  62                 _startTcs.TrySetResult(null);
  63             }
  64         }
  65
  66         public void Complete()
  67         {
  68             _channel.Writer.Complete();
  69         }
  70     }
  71 }
  72 #endif