diff --git a/pkg/stub/stub.go b/pkg/stub/stub.go index b0288bd6..56822224 100644 --- a/pkg/stub/stub.go +++ b/pkg/stub/stub.go @@ -283,6 +283,15 @@ func WithLogger(logger nrilog.Logger) Option { } } +// WithPluginRegistrationTimeout overrides the plugin's registration timeout +// (default DefaultRegistrationTimeout, 5s). +func WithPluginRegistrationTimeout(d time.Duration) Option { + return func(s *stub) error { + s.registrationTimeout = d + return nil + } +} + // stub implements Stub. type stub struct { sync.Mutex diff --git a/pkg/stub/stub_test.go b/pkg/stub/stub_test.go new file mode 100644 index 00000000..c1a550e8 --- /dev/null +++ b/pkg/stub/stub_test.go @@ -0,0 +1,54 @@ +/* + Copyright The containerd Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package stub_test + +import ( + "context" + "testing" + "time" + + "github.com/containerd/nri/pkg/api" + "github.com/containerd/nri/pkg/stub" + "github.com/stretchr/testify/require" +) + +type testPlugin struct{} + +func (testPlugin) RunPodSandbox(context.Context, *api.PodSandbox) error { + return nil +} + +func TestWithPluginRegistrationTimeout(t *testing.T) { + const want = 7 * time.Second + + s, err := stub.New(testPlugin{}, + stub.WithPluginName("test"), + stub.WithPluginIdx("00"), + stub.WithPluginRegistrationTimeout(want), + ) + require.NoError(t, err) + require.Equal(t, want, s.RegistrationTimeout()) +} + +func TestDefaultRegistrationTimeout(t *testing.T) { + s, err := stub.New(testPlugin{}, + stub.WithPluginName("test"), + stub.WithPluginIdx("00"), + ) + require.NoError(t, err) + require.Equal(t, stub.DefaultRegistrationTimeout, s.RegistrationTimeout()) +}