@@ -2,6 +2,7 @@ package k6
22
33import (
44 "context"
5+ grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
56 "net/http"
67 "net/http/httptest"
78 "runtime/pprof"
@@ -153,6 +154,83 @@ func TestLabelsFromBaggageUnaryInterceptor(t *testing.T) {
153154 })
154155}
155156
157+ func TestLabelsFromBaggageStreamInterceptor (t * testing.T ) {
158+ info := & grpc.StreamServerInfo {
159+ FullMethod : "/example.ExampleService/Test" ,
160+ }
161+
162+ t .Run ("adds_k6_labels_from_grpc_baggage" , func (t * testing.T ) {
163+ testCtx := testAddBaggageToGRPCRequest (t , context .Background (),
164+ "k6.test_run_id" , "123" ,
165+ "not_k6.some_other_key" , "value" ,
166+ )
167+
168+ handler := func (srv any , stream grpc.ServerStream ) error {
169+ ctx := stream .Context ()
170+
171+ b := baggage .FromContext (ctx )
172+ require .NotNil (t , b )
173+ testAssertEqualMembers (t , b .Members (),
174+ "k6.test_run_id" , "123" ,
175+ "not_k6.some_other_key" , "value" ,
176+ )
177+
178+ val , ok := pprof .Label (ctx , "k6_test_run_id" )
179+ require .True (t , ok )
180+ require .Equal (t , "123" , val )
181+
182+ _ , ok = pprof .Label (ctx , "not_k6_some_other_key" )
183+ require .False (t , ok )
184+
185+ return nil
186+ }
187+
188+ err := LabelsFromBaggageStreamInterceptor (nil , & grpcmiddleware.WrappedServerStream {WrappedContext : testCtx }, info , handler )
189+ require .NoError (t , err )
190+ })
191+
192+ t .Run ("passthrough_requests_with_no_baggage" , func (t * testing.T ) {
193+ testCtx := testAddBaggageToGRPCRequest (t , context .Background ())
194+
195+ handler := func (srv any , stream grpc.ServerStream ) error {
196+ ctx := stream .Context ()
197+
198+ b := baggage .FromContext (ctx )
199+ require .NotNil (t , b )
200+ require .Equal (t , 0 , b .Len ())
201+
202+ return nil
203+ }
204+
205+ err := LabelsFromBaggageStreamInterceptor (nil , & grpcmiddleware.WrappedServerStream {WrappedContext : testCtx }, info , handler )
206+ require .NoError (t , err )
207+ })
208+
209+ t .Run ("passthrough_requests_with_no_k6_baggage" , func (t * testing.T ) {
210+ testCtx := testAddBaggageToGRPCRequest (t , context .Background (),
211+ "not_k6.some_other_key" , "value" ,
212+ )
213+
214+ handler := func (srv any , stream grpc.ServerStream ) error {
215+ ctx := stream .Context ()
216+
217+ b := baggage .FromContext (ctx )
218+ require .NotNil (t , b )
219+ testAssertEqualMembers (t , b .Members (),
220+ "not_k6.some_other_key" , "value" ,
221+ )
222+
223+ _ , ok := pprof .Label (ctx , "not_k6_some_other_key" )
224+ require .False (t , ok )
225+
226+ return nil
227+ }
228+
229+ err := LabelsFromBaggageStreamInterceptor (nil , & grpcmiddleware.WrappedServerStream {WrappedContext : testCtx }, info , handler )
230+ require .NoError (t , err )
231+ })
232+ }
233+
156234func Test_setBaggageContextFromHeader (t * testing.T ) {
157235 t .Run ("sets_baggage_context" , func (t * testing.T ) {
158236 req := httptest .NewRequest ("GET" , "http://example.com" , nil )
0 commit comments