github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/io/mongodbio/helper_test.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 mongodbio
    17  
    18  import (
    19  	"context"
    20  	"testing"
    21  
    22  	"github.com/apache/beam/sdks/v2/go/test/integration/internal/containers"
    23  	"go.mongodb.org/mongo-driver/bson"
    24  	"go.mongodb.org/mongo-driver/bson/primitive"
    25  	"go.mongodb.org/mongo-driver/mongo"
    26  	"go.mongodb.org/mongo-driver/mongo/options"
    27  	"go.mongodb.org/mongo-driver/mongo/readpref"
    28  )
    29  
    30  const (
    31  	mongoImage = "mongo:6.0.3"
    32  	mongoPort  = "27017/tcp"
    33  	maxRetries = 5
    34  )
    35  
    36  func setUpTestContainer(ctx context.Context, t *testing.T) string {
    37  	t.Helper()
    38  
    39  	container := containers.NewContainer(
    40  		ctx,
    41  		t,
    42  		mongoImage,
    43  		maxRetries,
    44  		containers.WithPorts([]string{mongoPort}),
    45  	)
    46  
    47  	return containers.Port(ctx, t, container, mongoPort)
    48  }
    49  
    50  func objectIDFromHex(t *testing.T, hex string) primitive.ObjectID {
    51  	t.Helper()
    52  
    53  	id, err := primitive.ObjectIDFromHex(hex)
    54  	if err != nil {
    55  		t.Fatalf("error parsing hex string to primitive.ObjectID: %v", err)
    56  	}
    57  
    58  	return id
    59  }
    60  
    61  func newClient(ctx context.Context, t *testing.T, uri string) *mongo.Client {
    62  	t.Helper()
    63  
    64  	opts := options.Client().ApplyURI(uri)
    65  
    66  	client, err := mongo.Connect(ctx, opts)
    67  	if err != nil {
    68  		t.Fatalf("error connecting to MongoDB: %v", err)
    69  	}
    70  
    71  	t.Cleanup(func() {
    72  		if err := client.Disconnect(ctx); err != nil {
    73  			t.Fatalf("error disconnecting from MongoDB: %v", err)
    74  		}
    75  	})
    76  
    77  	if err := client.Ping(ctx, readpref.Primary()); err != nil {
    78  		t.Fatalf("error pinging MongoDB: %v", err)
    79  	}
    80  
    81  	return client
    82  }
    83  
    84  func dropCollection(ctx context.Context, t *testing.T, collection *mongo.Collection) {
    85  	t.Helper()
    86  
    87  	if err := collection.Drop(ctx); err != nil {
    88  		t.Fatalf("error dropping collection: %v", err)
    89  	}
    90  }
    91  
    92  func readDocuments(
    93  	ctx context.Context,
    94  	t *testing.T,
    95  	collection *mongo.Collection,
    96  ) []bson.M {
    97  	t.Helper()
    98  
    99  	cursor, err := collection.Find(ctx, bson.M{})
   100  	if err != nil {
   101  		t.Fatalf("error finding documents: %v", err)
   102  	}
   103  
   104  	var documents []bson.M
   105  	if err = cursor.All(ctx, &documents); err != nil {
   106  		t.Fatalf("error decoding documents: %v", err)
   107  	}
   108  
   109  	return documents
   110  }
   111  
   112  func writeDocuments(
   113  	ctx context.Context,
   114  	t *testing.T,
   115  	collection *mongo.Collection,
   116  	documents []any,
   117  ) {
   118  	t.Helper()
   119  
   120  	if _, err := collection.InsertMany(ctx, documents); err != nil {
   121  		t.Fatalf("error inserting documents: %v", err)
   122  	}
   123  }