github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/internal/containers/containers.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  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  
    16  // Package containers contains utilities for running test containers in integration tests.
    17  package containers
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/docker/go-connections/nat"
    25  	"github.com/testcontainers/testcontainers-go"
    26  	"github.com/testcontainers/testcontainers-go/wait"
    27  	"gopkg.in/retry.v1"
    28  )
    29  
    30  type ContainerOptionFn func(*testcontainers.ContainerRequest)
    31  
    32  func WithEnv(env map[string]string) ContainerOptionFn {
    33  	return func(option *testcontainers.ContainerRequest) {
    34  		option.Env = env
    35  	}
    36  }
    37  
    38  func WithHostname(hostname string) ContainerOptionFn {
    39  	return func(option *testcontainers.ContainerRequest) {
    40  		option.Hostname = hostname
    41  	}
    42  }
    43  
    44  func WithPorts(ports []string) ContainerOptionFn {
    45  	return func(option *testcontainers.ContainerRequest) {
    46  		option.ExposedPorts = ports
    47  	}
    48  }
    49  
    50  func WithWaitStrategy(waitStrategy wait.Strategy) ContainerOptionFn {
    51  	return func(option *testcontainers.ContainerRequest) {
    52  		option.WaitingFor = waitStrategy
    53  	}
    54  }
    55  
    56  func NewContainer(
    57  	ctx context.Context,
    58  	t *testing.T,
    59  	image string,
    60  	maxRetries int,
    61  	opts ...ContainerOptionFn,
    62  ) testcontainers.Container {
    63  	t.Helper()
    64  
    65  	request := testcontainers.ContainerRequest{Image: image}
    66  
    67  	for _, opt := range opts {
    68  		opt(&request)
    69  	}
    70  
    71  	genericRequest := testcontainers.GenericContainerRequest{
    72  		ContainerRequest: request,
    73  		Started:          true,
    74  	}
    75  
    76  	strategy := retry.LimitCount(
    77  		maxRetries,
    78  		retry.Exponential{
    79  			Initial: time.Second,
    80  			Factor:  2,
    81  		},
    82  	)
    83  
    84  	var container testcontainers.Container
    85  	var err error
    86  
    87  	for attempt := retry.Start(strategy, nil); attempt.Next(); {
    88  		container, err = testcontainers.GenericContainer(ctx, genericRequest)
    89  		if err == nil {
    90  			break
    91  		}
    92  
    93  		if attempt.Count() == maxRetries {
    94  			t.Fatalf("failed to start container with %v retries: %v", maxRetries, err)
    95  		}
    96  	}
    97  
    98  	t.Cleanup(func() {
    99  		if err := container.Terminate(ctx); err != nil {
   100  			t.Fatalf("error terminating container: %v", err)
   101  		}
   102  	})
   103  
   104  	return container
   105  }
   106  
   107  func Port(
   108  	ctx context.Context,
   109  	t *testing.T,
   110  	container testcontainers.Container,
   111  	port nat.Port,
   112  ) string {
   113  	t.Helper()
   114  
   115  	mappedPort, err := container.MappedPort(ctx, port)
   116  	if err != nil {
   117  		t.Fatalf("error getting mapped port: %v", err)
   118  	}
   119  
   120  	return mappedPort.Port()
   121  }