SERVER-409050 auth test for searchBeta agg stage

2 files changed, 73 insertions, 8 deletions

diff --git a/jstests/auth/lib/commands_lib.js b/jstests/auth/lib/commands_lib.js
index 96970cdef91..2580ebe6635 100644
--- a/jstests/auth/lib/commands_lib.js
+++ b/jstests/auth/lib/commands_lib.js
@@ -5933,6 +5933,46 @@ var authCommandsLib = {
           }
         },
         {
+          testname: "aggregate_$searchBeta",
+          command: {
+              aggregate: "foo",
+              cursor: {},
+              pipeline: [{
+                  $searchBeta: {
+                      // empty query
+                  }
+              }]
+          },
+          skipSharded: false,
+          // Only enterprise knows of this aggregation stage.
+          skipTest:
+              (conn) =>
+                  !conn.getDB("admin").runCommand({buildInfo: 1}).modules.includes("enterprise"),
+          testcases: [
+              {
+                runOnDb: firstDbName,
+                roles: roles_read,
+                privileges: [{resource: {db: firstDbName, collection: "foo"}, actions: ["find"]}]
+              },
+              {
+                runOnDb: secondDbName,
+                roles: roles_readAny,
+                privileges:
+                    [{resource: {db: secondDbName, collection: "foo"}, actions: ["find"]}]
+              }
+          ],
+          setup: function(db) {
+              // Configure the $searchBeta stage to always return EOF so we can avoid the hassle
+              // of giving mongod a host and port for mongot.
+              const cmd = {configureFailPoint: "searchBetaReturnEofImmediately", mode: "alwaysOn"};
+              FixtureHelpers.runCommandOnEachPrimary({db: db.getSiblingDB("admin"), cmdObj: cmd});
+          },
+          teardown: function(db) {
+              const cmd = {configureFailPoint: "searchBetaReturnEofImmediately", mode: "off"};
+              FixtureHelpers.runCommandOnEachPrimary({db: db.getSiblingDB("admin"), cmdObj: cmd});
+          }
+        },
+        {
           testname: "startRecordingTraffic",
           command: {startRecordingTraffic: 1, filename: "notARealPath"},
           testcases: [
@@ -5979,18 +6019,18 @@ var authCommandsLib = {
      *  An array of strings. Each string in the array reports
      *  a particular test error.
      */
-    runOneTest: function(conn, t, impls) {
+    runOneTest: function(conn, t, impls, isMongos) {
         jsTest.log("Running test: " + t.testname);
 
         if (t.skipTest && t.skipTest(conn)) {
             return [];
         }
         // some tests shouldn't run in a sharded environment
-        if (t.skipSharded && this.isMongos(conn)) {
+        if (t.skipSharded && isMongos) {
             return [];
         }
         // others shouldn't run in a standalone environment
-        if (t.skipUnlessSharded && !this.isMongos(conn)) {
+        if (t.skipUnlessSharded && !isMongos) {
             return [];
         }
         // some tests require replica sets to be enabled.
@@ -6043,8 +6083,9 @@ var authCommandsLib = {
 
         var failures = [];
 
+        const isMongos = this.isMongos(conn);
         for (var i = 0; i < this.tests.length; i++) {
-            res = this.runOneTest(conn, this.tests[i], impls);
+            res = this.runOneTest(conn, this.tests[i], impls, isMongos);
             failures = failures.concat(res);
         }
 
diff --git a/jstests/libs/fixture_helpers.js b/jstests/libs/fixture_helpers.js
index a1b0bbbb968..dc4fc83365b 100644
--- a/jstests/libs/fixture_helpers.js
+++ b/jstests/libs/fixture_helpers.js
@@ -84,12 +84,36 @@ var FixtureHelpers = (function() {
      * Runs the command given by 'cmdObj' on the database given by 'db' on each replica set in
      * the fixture (besides the config servers). Asserts that each command works, and returns an
      * array with the responses from each shard, or with a single element if the fixture was a
-     * replica set. Asserts if the fixture is a standalone or if the shards are standalones.
+     * replica set. If the fixture is a standalone, will run the command directly.
      */
     function runCommandOnEachPrimary({db, cmdObj}) {
-        return _getAllReplicas(db).map(
-            (replSet) =>
-                assert.commandWorked(replSet.getPrimary().getDB(db.getName()).runCommand(cmdObj)));
+        function getConnToPrimaryOrStandalone(host) {
+            const conn = new Mongo(host);
+            const isMaster = conn.getDB("test").isMaster();
+
+            if (isMaster.hasOwnProperty("setName")) {
+                // It's a repl set.
+                const rs = new ReplSetTest(host);
+                return rs.getPrimary();
+            } else {
+                // It's a standalone.
+                return conn;
+            }
+        }
+
+        const connList = [];
+        if (isMongos(db)) {
+            const shardObjs = db.getSiblingDB("config").shards.find().sort({_id: 1}).toArray();
+
+            for (let shardObj of shardObjs) {
+                connList.push(getConnToPrimaryOrStandalone(shardObj.host));
+            }
+        } else {
+            connList.push(getConnToPrimaryOrStandalone(db.getMongo().host));
+        }
+
+        return connList.map((conn) =>
+                                assert.commandWorked(conn.getDB(db.getName()).runCommand(cmdObj)));
     }
 
     /**