SERVER-50157 Add Sequence to network mock framework

3 files changed, 226 insertions, 1 deletions

diff --git a/src/mongo/executor/mock_network_fixture.cpp b/src/mongo/executor/mock_network_fixture.cpp
index da028d8a6ba..766397bd154 100644
--- a/src/mongo/executor/mock_network_fixture.cpp
+++ b/src/mongo/executor/mock_network_fixture.cpp
@@ -95,7 +95,7 @@ void MockNetwork::_runUntilIdle() {
             // (Iterating a vector backwards is much cheaper than pushing to its front.)
             auto const& exp =
                 std::find_if(_expectations.rbegin(), _expectations.rend(), [&](const auto& exp) {
-                    return !exp->isSatisfied() && exp->match(request);
+                    return !exp->isSatisfied() && exp->prerequisitesMet() && exp->match(request);
                 });
 
             if (exp != _expectations.rend()) {
diff --git a/src/mongo/executor/mock_network_fixture.h b/src/mongo/executor/mock_network_fixture.h
index a2628042097..91c08207898 100644
--- a/src/mongo/executor/mock_network_fixture.h
+++ b/src/mongo/executor/mock_network_fixture.h
@@ -99,6 +99,14 @@ public:
             : _matcher(std::move(matcher)), _action(std::move(action)) {}
         virtual ~Expectation() {}
 
+        // We should only try to match an expectation if all its prerequisites in the sequence are
+        // satisfied (or if it is not part of a sequence in the first place).
+        // Default expectations should always match, as they cannot form sequences, and because we
+        // generally have no restrictions on when and how many times they match.
+        virtual bool prerequisitesMet() {
+            return true;
+        }
+
         bool match(const BSONObj& request) {
             return _matcher(request);
         }
@@ -136,6 +144,17 @@ public:
             _allowedTimes = 1;
         }
 
+        // We forbid matching if we have unsatisfied prerequisites and allow it otherwise.
+        bool prerequisitesMet() override {
+            // No prerequisites at all - okay to match.
+            if (!_prevInSequence) {
+                return true;
+            }
+            // If this expectation's immediate prerequisite is satisfied, that implies the
+            // previous ones in the chain have also been satisfied.
+            return _prevInSequence->isSatisfied();
+        }
+
         Expectation& times(int t) {
             _allowedTimes = t;
             return *this;
@@ -148,6 +167,9 @@ public:
         void checkSatisfied() override {
             uassert(5015501, "UserExpectation not satisfied", isSatisfied());
         }
+
+        // The expectation that comes before this on in the sequence, if applicable.
+        UserExpectation* _prevInSequence = nullptr;
     };
 
     class DefaultExpectation : public Expectation {
@@ -159,6 +181,32 @@ public:
         }
     };
 
+    class Sequence {
+    public:
+        void addExpectation(UserExpectation& exp) {
+            exp._prevInSequence = _lastExpectation;
+            _lastExpectation = &exp;
+        }
+
+    private:
+        UserExpectation* _lastExpectation = nullptr;
+    };
+
+    // RAII construct to handle expectations declared as part of a sequence.
+    class InSequence {
+    public:
+        InSequence(MockNetwork& mock) : _mock(mock) {
+            invariant(!_mock._activeSequence);
+            _mock._activeSequence = std::make_unique<Sequence>();
+        };
+        ~InSequence() {
+            _mock._activeSequence.reset();
+        };
+
+    private:
+        MockNetwork& _mock;
+    };
+
     explicit MockNetwork(executor::NetworkInterfaceMock* net) : _net(net) {}
 
     // Accept anything that Matcher's and Action's constructors allow.
@@ -170,6 +218,11 @@ public:
                                                      std::move(action));
         auto& ref = *exp;
         _expectations.emplace_back(std::move(exp));
+
+        if (_activeSequence) {
+            _activeSequence.get()->addExpectation(ref);
+        }
+
         return ref;
     }
 
@@ -229,6 +282,7 @@ private:
     bool _allExpectationsSatisfied() const;
 
     std::vector<std::unique_ptr<Expectation>> _expectations;
+    std::unique_ptr<Sequence> _activeSequence;
     executor::NetworkInterfaceMock* _net;
 };
 
diff --git a/src/mongo/executor/mock_network_fixture_test.cpp b/src/mongo/executor/mock_network_fixture_test.cpp
index caf12ad2cc1..86ae5122b60 100644
--- a/src/mongo/executor/mock_network_fixture_test.cpp
+++ b/src/mongo/executor/mock_network_fixture_test.cpp
@@ -84,11 +84,17 @@ public:
 
     std::string kExampleCmdName = "someCommandName";
     std::string kExampleCmdNameTwo = kExampleCmdName + "_two";
+    std::string kExampleCmdNameThree = kExampleCmdName + "_three";
+    std::string kExampleCmdNameFour = kExampleCmdName + "_four";
 
     RemoteCommandRequestOnAny kExampleRequest{
         {testHost()}, "testDB", BSON(kExampleCmdName << 1), rpc::makeEmptyMetadata(), nullptr};
     RemoteCommandRequestOnAny kExampleRequestTwo{
         {testHost()}, "testDB", BSON(kExampleCmdNameTwo << 1), rpc::makeEmptyMetadata(), nullptr};
+    RemoteCommandRequestOnAny kExampleRequestThree{
+        {testHost()}, "testDB", BSON(kExampleCmdNameThree << 1), rpc::makeEmptyMetadata(), nullptr};
+    RemoteCommandRequestOnAny kExampleRequestFour{
+        {testHost()}, "testDB", BSON(kExampleCmdNameFour << 1), rpc::makeEmptyMetadata(), nullptr};
 
     BSONObj kExampleResponse = BSON("some"
                                     << "response");
@@ -294,5 +300,170 @@ TEST_F(MockNetworkTest, MockFixtureNotEnoughTimesMatched) {
     evaluateResponses(2 /* numExpected */);
 }
 
+TEST_F(MockNetworkTest, MockFixtureSequenceTestCorrectOrder) {
+    // The requests will come in the expected order.
+    {
+        MockNetwork::InSequence seq(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+    }
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestOne));
+    ASSERT_OK(startCommand(requestTwo));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(2 /* numExpected */);
+}
+
+TEST_F(MockNetworkTest, MockFixtureSequenceTestOppositeOrder) {
+    // We expect One -> Two, but the requests will come in the opposite order.
+    // We will refuse to match Two to its user expectation if One is not satisfied,
+    // so we make a default expectation to fall back to so the test can move forward.
+    mock().defaultExpect(kExampleCmdNameTwo, kExampleResponse);
+
+    {
+        MockNetwork::InSequence seq(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+    }
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestTwo));
+    ASSERT_OK(startCommand(requestOne));
+
+    mock().runUntil(net().now() + Milliseconds(100));
+
+    // The second expectation will not have been satisfied, as the first one has not yet been
+    // met at that time it came in.
+    ASSERT_THROWS_CODE(mock().verifyExpectations(), DBException, (ErrorCodes::Error)5015501);
+
+    // Now we can send a request for expectation Two again.
+    ASSERT_OK(startCommand(requestTwo));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(3 /* numExpected */);
+}
+
+TEST_F(MockNetworkTest, MockFixtureSequenceTestCorrectPartialOrder) {
+    // We only require One -> Two, so One -> Three -> Two is valid.
+    {
+        MockNetwork::InSequence seq(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+    }
+
+    mock().expect(kExampleCmdNameThree, kExampleResponse).times(1);
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+    RemoteCommandRequestOnAny requestThree{kExampleRequestThree};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestOne));
+    ASSERT_OK(startCommand(requestThree));
+    ASSERT_OK(startCommand(requestTwo));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(3 /* numExpected */);
+}
+
+TEST_F(MockNetworkTest, MockFixtureSequenceTestTwoSequencesNoOverlap) {
+    // We have One -> Two and Three -> Four. We run One -> Two -> Three -> Four;
+    {
+        MockNetwork::InSequence seqOne(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+    }
+
+    {
+        MockNetwork::InSequence seqTwo(mock());
+
+        mock().expect(kExampleCmdNameThree, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameFour, kExampleResponse).times(1);
+    }
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+    RemoteCommandRequestOnAny requestThree{kExampleRequestThree};
+    RemoteCommandRequestOnAny requestFour{kExampleRequestFour};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestOne));
+    ASSERT_OK(startCommand(requestTwo));
+    ASSERT_OK(startCommand(requestThree));
+    ASSERT_OK(startCommand(requestFour));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(4 /* numExpected */);
+}
+
+TEST_F(MockNetworkTest, MockFixtureSequenceTestTwoSequencesPartialOrder) {
+    // We have One -> Two and Three -> Four. We run One -> Three -> Two -> Four;
+    {
+        MockNetwork::InSequence seqOne(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+    }
+
+    {
+        MockNetwork::InSequence seqTwo(mock());
+
+        mock().expect(kExampleCmdNameThree, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameFour, kExampleResponse).times(1);
+    }
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+    RemoteCommandRequestOnAny requestThree{kExampleRequestThree};
+    RemoteCommandRequestOnAny requestFour{kExampleRequestFour};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestOne));
+    ASSERT_OK(startCommand(requestThree));
+    ASSERT_OK(startCommand(requestTwo));
+    ASSERT_OK(startCommand(requestFour));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(4 /* numExpected */);
+}
+
+TEST_F(MockNetworkTest, MockFixtureSequenceTestLongerChain) {
+    // One -> Two -> Three -> Four required.
+    {
+        MockNetwork::InSequence seqOne(mock());
+
+        mock().expect(kExampleCmdName, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameTwo, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameThree, kExampleResponse).times(1);
+        mock().expect(kExampleCmdNameFour, kExampleResponse).times(1);
+    }
+
+    RemoteCommandRequestOnAny requestOne{kExampleRequest};
+    RemoteCommandRequestOnAny requestTwo{kExampleRequestTwo};
+    RemoteCommandRequestOnAny requestThree{kExampleRequestThree};
+    RemoteCommandRequestOnAny requestFour{kExampleRequestFour};
+
+    // Run commands in this specific order.
+    ASSERT_OK(startCommand(requestOne));
+    ASSERT_OK(startCommand(requestTwo));
+    ASSERT_OK(startCommand(requestThree));
+    ASSERT_OK(startCommand(requestFour));
+
+    mock().runUntilExpectationsSatisfied();
+    evaluateResponses(4 /* numExpected */);
+}
+
 }  // namespace
 }  // namespace mongo
\ No newline at end of file