From a5e3007d307867d874a9377bcd172b88df8163f7 Mon Sep 17 00:00:00 2001
From: Bernard Gorman <bernard.gorman@gmail.com>
Date: Sun, 12 Nov 2017 19:13:34 +0000
Subject: SERVER-31394 Create passthrough of existing $changeStream tests to
 run against sharded collections

---
 jstests/libs/change_stream_util.js                 | 13 ++++++----
 jstests/libs/collection_drop_recreate.js           | 29 ++++++++++++++++++++++
 .../implicitly_shard_accessed_collections.js       | 10 +++++---
 3 files changed, 43 insertions(+), 9 deletions(-)
 create mode 100644 jstests/libs/collection_drop_recreate.js

(limited to 'jstests/libs')

diff --git a/jstests/libs/change_stream_util.js b/jstests/libs/change_stream_util.js
index cbaf69fdedf..aea8aa3a1a5 100644
--- a/jstests/libs/change_stream_util.js
+++ b/jstests/libs/change_stream_util.js
@@ -30,10 +30,14 @@ function ChangeStreamTest(_db, name = "ChangeStreamTest") {
             pipeline.push(self.oplogProjection);
         }
 
+        // The 'collection' argument may be either a collection name or DBCollection object.
+        assert(collection instanceof DBCollection || typeof collection === "string");
+        const collName = (collection instanceof DBCollection ? collection.getName() : collection);
+
         let res = assert.commandWorked(_db.runCommand(
-            Object.merge({aggregate: collection.getName(), pipeline: pipeline}, aggregateOptions)));
+            Object.merge({aggregate: collName, pipeline: pipeline}, aggregateOptions)));
         assert.neq(res.cursor.id, 0);
-        _allCursors.push({db: _db.getName(), coll: collection.getName(), cursorId: res.cursor.id});
+        _allCursors.push({db: _db.getName(), coll: collName, cursorId: res.cursor.id});
         return res.cursor;
     };
 
@@ -41,10 +45,9 @@ function ChangeStreamTest(_db, name = "ChangeStreamTest") {
      * Issues a 'getMore' on the provided cursor and returns the cursor returned.
      */
     self.getNextBatch = function(cursor) {
-        let collection = _db.getMongo().getCollection(cursor.ns);
+        const collName = cursor.ns.split(/\.(.+)/)[1];
         return assert
-            .commandWorked(_db.runCommand(
-                {getMore: cursor.id, collection: collection.getName(), batchSize: 1}))
+            .commandWorked(_db.runCommand({getMore: cursor.id, collection: collName, batchSize: 1}))
             .cursor;
     };
 
diff --git a/jstests/libs/collection_drop_recreate.js b/jstests/libs/collection_drop_recreate.js
new file mode 100644
index 00000000000..9d2888ce6d8
--- /dev/null
+++ b/jstests/libs/collection_drop_recreate.js
@@ -0,0 +1,29 @@
+/**
+ * Attempts to drop the given collection and asserts either that the drop succeeded or the
+ * collection did not exist. Avoids automatically recreating the collection in the case of test
+ * suites where accessing or dropping the collection implicitly recreates it.
+ */
+function assertDropCollection(db, collName) {
+    var cmdRes = db.runCommand({drop: collName});
+    assert(cmdRes.ok === 1 || cmdRes.code === ErrorCodes.NamespaceNotFound, tojson(cmdRes));
+}
+
+/**
+ * Attempts to create a collection with the given name and options, if any, and asserts on failure.
+ * Returns the newly-created collection on success. When running under a sharded collections
+ * passthrough, the new collection will be implicitly sharded.
+ */
+function assertCreateCollection(db, collName, collOpts) {
+    assert.commandWorked(db.createCollection(collName, collOpts));
+    return db.getCollection(collName);
+}
+
+/**
+ * Attempts to drop a collection with the given name and recreate it with the specified options, if
+ * any. Asserts if either step fails. Returns the newly-created collection on success. When running
+ * under a sharded collections passthrough, the new collection will be implicitly sharded.
+ */
+function assertDropAndRecreateCollection(db, collName, collOpts) {
+    assertDropCollection(db, collName);
+    return assertCreateCollection(db, collName, collOpts);
+}
\ No newline at end of file
diff --git a/jstests/libs/override_methods/implicitly_shard_accessed_collections.js b/jstests/libs/override_methods/implicitly_shard_accessed_collections.js
index 3f1f8b24a55..e8ebf477167 100644
--- a/jstests/libs/override_methods/implicitly_shard_accessed_collections.js
+++ b/jstests/libs/override_methods/implicitly_shard_accessed_collections.js
@@ -45,7 +45,8 @@
             assert.commandWorked(res, "enabling sharding on the '" + dbName + "' db failed");
         }
 
-        res = db.adminCommand({shardCollection: fullName, key: {_id: 'hashed'}});
+        res = db.adminCommand(
+            {shardCollection: fullName, key: {_id: 'hashed'}, collation: {locale: "simple"}});
         if (res.ok === 0 && testMayRunDropInParallel) {
             // We ignore ConflictingOperationInProgress error responses from the
             // "shardCollection" command if it's possible the test was running a "drop" command
@@ -64,9 +65,10 @@
     DB.prototype.getCollection = function() {
         var collection = originalGetCollection.apply(this, arguments);
 
-        // If the collection exists, there must have been a previous call to getCollection
-        // where we sharded the collection so there's no need to do it again.
-        if (collection.exists()) {
+        const collStats = this.runCommand({collStats: collection.getName()});
+
+        // If the collection is already sharded or is non-empty, do not attempt to shard.
+        if (collStats.sharded || collStats.count > 0) {
             return collection;
         }
 
-- 
cgit v1.2.1