diff --git a/main.go b/main.go
index 7bcf2ac..c28793d 100644
--- a/main.go
+++ b/main.go
@@ -37,11 +37,11 @@ func main() {
 
 	collection := client.Database("benchmarking").Collection("testdata")
 
-	if testType == "insert" {
+	if testType == "insert" || testType == "upsert" {
 		if err := collection.Drop(context.Background()); err != nil {
 			log.Fatalf("Failed to drop collection: %v", err)
 		}
-		log.Println("Collection dropped. Starting new insert test...")
+		log.Println("Collection dropped. Starting new rate test...")
 	} else {
 		log.Printf("Starting %s test...\n", testType)
 	}
@@ -140,7 +140,7 @@ func main() {
 					}
 
 				case "upsert":
-					docID := int64(rand.Intn(docCount))
+					docID := int64(rand.Intn(docCount / 2))
 					filter := bson.M{"_id": docID}
 					update := bson.M{"$set": bson.M{"updatedAt": time.Now().Unix(), "rnd": rand.Int63()}}
 					opts := options.Update().SetUpsert(true)
diff --git a/mongo-bench-entrypoint.sh b/mongo-bench-entrypoint.sh
index 7604a9c..350b9d2 100755
--- a/mongo-bench-entrypoint.sh
+++ b/mongo-bench-entrypoint.sh
@@ -31,6 +31,17 @@ DOC_COUNT=$(mongosh 'mongodb://root:example@mongodb:27017/?authSource=admin' --q
 if [ "$DOC_COUNT" -ne 0 ]; then
   echo "Error: Expected 0 documents, found $DOC_COUNT"
   exit 1
-else
+fi
+
+# Run the upsert test
+echo 'Running upsert test...'
+./mongo-bench --uri mongodb://root:example@mongodb:27017 --type upsert --threads 10 --docs 80000
+
+echo 'Checking document count...'
+DOC_COUNT=$(mongosh 'mongodb://root:example@mongodb:27017/?authSource=admin' --quiet --eval 'JSON.stringify({count: db.getSiblingDB("benchmarking").testdata.countDocuments()})' | jq -r '.count')
+if [ "$DOC_COUNT" -gt 0 ]; then
   echo 'All tests passed successfully.'
+else
+  echo "Error: Expected >0 documents, found $DOC_COUNT"
+  exit 1
 fi
\ No newline at end of file