content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.instrumentation.mongo.v3_1;
  
  import static java.util.Arrays.asList;
  import static org.assertj.core.api.Assertions.assertThat;
  import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
  
  import com.mongodb.MongoClient;
  import com.mongodb.MongoClientOptions;
  import com.mongodb.MongoTimeoutException;
  import com.mongodb.ServerAddress;
  import com.mongodb.client.MongoCollection;
  import com.mongodb.client.MongoDatabase;
  import com.mongodb.client.result.DeleteResult;
  import com.mongodb.client.result.UpdateResult;
  import io.opentelemetry.instrumentation.mongo.testing.AbstractMongoClientTest;
  import io.opentelemetry.instrumentation.test.utils.PortUtils;
  import java.util.ArrayList;
  import org.bson.BsonDocument;
  import org.bson.BsonString;
  import org.bson.Document;
  import org.junit.jupiter.api.AfterAll;
  import org.junit.jupiter.api.BeforeAll;
  import org.junit.jupiter.api.Test;
  
  public abstract class AbstractMongo31ClientTest
      extends AbstractMongoClientTest<MongoCollection<Document>> {
  
    protected abstract void configureMongoClientOptions(MongoClientOptions.Builder options);
  
    private MongoClient client;
  
    @BeforeAll
    void setup() {
      MongoClientOptions.Builder options =
          MongoClientOptions.builder().description("some-description");
      configureMongoClientOptions(options);
      client = new MongoClient(new ServerAddress(host, port), options.build());
    }
  
    @AfterAll
    void cleanup() {
      if (client != null) {
        client.close();
      }
    }
  
    @Override
    protected void createCollection(String dbName, String collectionName) {
      MongoDatabase db = client.getDatabase(dbName);
      db.createCollection(collectionName);
    }
  
    @Override
    protected void createCollectionNoDescription(String dbName, String collectionName) {
      MongoClientOptions.Builder options = MongoClientOptions.builder();
      configureMongoClientOptions(options);
      MongoDatabase db =
          new MongoClient(new ServerAddress(host, port), options.build()).getDatabase(dbName);
      db.createCollection(collectionName);
    }
  
    @Override
    protected void createCollectionWithAlreadyBuiltClientOptions(
        String dbName, String collectionName) {
      MongoClientOptions clientOptions = client.getMongoClientOptions();
      MongoClientOptions newClientOptions = MongoClientOptions.builder(clientOptions).build();
      MongoDatabase db =
          new MongoClient(new ServerAddress(host, port), newClientOptions).getDatabase(dbName);
      db.createCollection(collectionName);
    }
  
    @Override
    protected void createCollectionCallingBuildTwice(String dbName, String collectionName) {
      MongoClientOptions.Builder options =
          MongoClientOptions.builder().description("some-description");
      configureMongoClientOptions(options);
      options.build();
      MongoDatabase db =
          new MongoClient(new ServerAddress(host, port), options.build()).getDatabase(dbName);
      db.createCollection(collectionName);
    }
  
    @Override
    protected long getCollection(String dbName, String collectionName) {
      MongoDatabase db = client.getDatabase(dbName);
      return db.getCollection(collectionName).count();
    }
  
    @Override
    protected MongoCollection<Document> setupInsert(String dbName, String collectionName) {
      MongoCollection<Document> collection =
          testing()
              .runWithSpan(
                  "setup",
                  () -> {
                    MongoDatabase db = client.getDatabase(dbName);
                    db.createCollection(collectionName);
                    return db.getCollection(collectionName);
                  });
      ignoreTracesAndClear(1);
      return collection;
    }
  
    @Override
    protected long insert(MongoCollection<Document> collection) {
      collection.insertOne(new Document("password", "SECRET"));
      return collection.count();
    }
  
    @Override
    protected MongoCollection<Document> setupUpdate(String dbName, String collectionName) {
      MongoCollection<Document> collection =
          testing()
              .runWithSpan(
                  "setup",
                  () -> {
                    MongoDatabase db = client.getDatabase(dbName);
                    db.createCollection(collectionName);
                    MongoCollection<Document> coll = db.getCollection(collectionName);
                    coll.insertOne(new Document("password", "OLDPW"));
                    return coll;
                  });
      ignoreTracesAndClear(1);
      return collection;
    }
  
    @Override
    protected long update(MongoCollection<Document> collection) {
      UpdateResult result =
          collection.updateOne(
              new BsonDocument("password", new BsonString("OLDPW")),
              new BsonDocument("$set", new BsonDocument("password", new BsonString("NEWPW"))));
      collection.count();
      return result.getModifiedCount();
    }
  
    @Override
    protected MongoCollection<Document> setupDelete(String dbName, String collectionName) {
      MongoCollection<Document> collection =
          testing()
              .runWithSpan(
                  "setup",
                  () -> {
                    MongoDatabase db = client.getDatabase(dbName);
                    db.createCollection(collectionName);
                    MongoCollection<Document> coll = db.getCollection(collectionName);
                    coll.insertOne(new Document("password", "SECRET"));
                    return coll;
                  });
      ignoreTracesAndClear(1);
      return collection;
    }
  
    @Override
    protected long delete(MongoCollection<Document> collection) {
      DeleteResult result =
          collection.deleteOne(new BsonDocument("password", new BsonString("SECRET")));
      collection.count();
      return result.getDeletedCount();
    }
  
    @Override
    protected MongoCollection<Document> setupGetMore(String dbName, String collectionName) {
      MongoCollection<Document> collection =
          testing()
              .runWithSpan(
                  "setup",
                  () -> {
                    MongoDatabase db = client.getDatabase(dbName);
                    MongoCollection<Document> coll = db.getCollection(collectionName);
                    coll.insertMany(
                        asList(
                            new Document("_id", 0), new Document("_id", 1), new Document("_id", 2)));
                    return coll;
                  });
      ignoreTracesAndClear(1);
      return collection;
    }
  
    @Override
    protected void getMore(MongoCollection<Document> collection) {
      collection
          .find()
          .filter(new Document("_id", new Document("$gte", 0)))
          .batchSize(2)
          .into(new ArrayList<>());
    }
  
    @Override
    protected void error(String dbName, String collectionName) {
      MongoCollection<Document> collection =
          testing()
              .runWithSpan(
                  "setup",
                  () -> {
                    MongoDatabase db = client.getDatabase(dbName);
                    db.createCollection(collectionName);
                    return db.getCollection(collectionName);
                  });
      ignoreTracesAndClear(1);
      collection.updateOne(new BsonDocument(), new BsonDocument());
    }
  
    @Test
    void testClientFailure() {
      MongoClientOptions options = MongoClientOptions.builder().serverSelectionTimeout(10).build();
      MongoClient client = new MongoClient(new ServerAddress(host, PortUtils.UNUSABLE_PORT), options);
  
      assertThatExceptionOfType(MongoTimeoutException.class)
          .isThrownBy(
              () -> {
                MongoDatabase db = client.getDatabase("test_db");
                db.createCollection(createCollectionName());
              });
      // Unfortunately not caught by our instrumentation.
      assertThat(testing().spans()).isEmpty();
    }
  }
  