From 8a383f5516a5216a76d5db8bbaf504f2c6502f56 Mon Sep 17 00:00:00 2001
From: Harry Pidcock <harry.pidcock@canonical.com>
Date: Fri, 29 Jul 2022 16:17:59 +1000
Subject: [PATCH] Add started and asserted hooks for sstxn testing.

---
 .gitignore  |   1 +
 txn.go      |  30 ++++++++++++++
 txn_test.go | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 144 insertions(+)

diff --git a/.gitignore b/.gitignore
index 61ead86..aeb1259 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 /vendor
+.vscode/
diff --git a/txn.go b/txn.go
index 8bbe80e..c86cfe3 100755
--- a/txn.go
+++ b/txn.go
@@ -395,6 +395,32 @@ func (tr *transactionRunner) RunTransaction(transaction *Transaction) error {
 	}
 	start := tr.clock.Now()
 	runner := tr.newRunner()
+	if len(testHooks) > 0 {
+		if sstxnRunner, ok := runner.(*sstxn.Runner); ok {
+			started := testHooks[0].Started
+			if started != nil {
+				sstxnRunner.SetStartHook(func() {
+					logger.Infof("transaction 'started' hook start")
+					started()
+					logger.Infof("transaction 'started' hook end")
+				})
+			}
+			asserted := testHooks[0].Asserted
+			if asserted != nil {
+				sstxnRunner.SetPostAssertHook(func() {
+					logger.Infof("transaction 'asserted' hook start")
+					asserted()
+					logger.Infof("transaction 'asserted' hook end")
+				})
+			}
+			if started != nil || asserted != nil {
+				defer func() {
+					sstxnRunner.SetStartHook(nil)
+					sstxnRunner.SetPostAssertHook(nil)
+				}()
+			}
+		}
+	}
 	err := runner.Run(transaction.Ops, "", nil)
 	if tr.runTransactionObserver != nil {
 		transaction.Error = err
@@ -421,6 +447,10 @@ func (tr *transactionRunner) MaybePruneTransactions(pruneOpts PruneOptions) erro
 type TestHook struct {
 	Before func()
 	After  func()
+	// Started only works for serverside txns
+	Started func()
+	// Asserted only works for serverside txns
+	Asserted func()
 }
 
 // TestHooks returns the test hooks for a transaction runner.
diff --git a/txn_test.go b/txn_test.go
index 923b9e4..5ece369 100644
--- a/txn_test.go
+++ b/txn_test.go
@@ -264,6 +264,119 @@ func (s *txnSuite) TestBeforeHooks(c *gc.C) {
 	c.Assert(found, gc.DeepEquals, doc)
 }
 
+func (s *sstxnSuite) TestStartedHooks(c *gc.C) {
+	s.insertDoc(c, "1", "Simple")
+
+	secondSession := s.collection.Database.Session.Copy()
+	defer secondSession.Close()
+	secondRunner := jujutxn.NewRunner(jujutxn.RunnerParams{
+		Database:               secondSession.DB("juju"),
+		ChangeLogName:          "txns.log",
+		ServerSideTransactions: true,
+		PauseFunc: func(dur time.Duration) {
+			s.backoffs = append(s.backoffs, dur)
+		},
+	})
+
+	setDocName := func(id, name string) {
+		ops := []txn.Op{{
+			C:      s.collection.Name,
+			Id:     id,
+			Assert: txn.DocExists,
+			Update: bson.D{{"$set", bson.D{{"name", name}}}},
+		}}
+		err := secondRunner.RunTransaction(&jujutxn.Transaction{
+			Ops:     ops,
+			Attempt: 0,
+		})
+		c.Assert(err, gc.IsNil)
+	}
+
+	hooks := []jujutxn.TestHook{
+		{Started: func() { setDocName("1", "FooBar") }},
+		{Started: func() { setDocName("1", "Foo") }},
+	}
+	defer txntesting.SetTestHooks(c, s.txnRunner, hooks...).Check()
+	maxAttempt := 0
+	buildTxn := func(attempt int) ([]txn.Op, error) {
+		maxAttempt = attempt
+		ops := []txn.Op{{
+			C:      s.collection.Name,
+			Id:     "1",
+			Assert: bson.D{{"name", "Foo"}},
+			Update: bson.D{{"$set", bson.D{{"name", "Bar"}}}},
+		}}
+		return ops, nil
+	}
+	err := s.txnRunner.Run(buildTxn)
+	c.Assert(err, gc.IsNil)
+	var found simpleDoc
+	err = s.collection.FindId("1").One(&found)
+	c.Assert(err, gc.IsNil)
+	c.Assert(maxAttempt, gc.Equals, 1)
+	doc := simpleDoc{"1", "Bar"}
+	c.Assert(found, gc.DeepEquals, doc)
+}
+
+func (s *sstxnSuite) TestAssertedHooks(c *gc.C) {
+	s.insertDoc(c, "1", "Foo")
+
+	secondSession := s.collection.Database.Session.Copy()
+	defer secondSession.Close()
+	secondRunner := jujutxn.NewRunner(jujutxn.RunnerParams{
+		Database:               secondSession.DB("juju"),
+		ChangeLogName:          "txns.log",
+		ServerSideTransactions: true,
+		PauseFunc: func(dur time.Duration) {
+			s.backoffs = append(s.backoffs, dur)
+		},
+	})
+
+	setDocName := func(id, name string) {
+		ops := []txn.Op{{
+			C:      s.collection.Name,
+			Id:     id,
+			Assert: txn.DocExists,
+			Update: bson.D{{"$set", bson.D{{"name", name}}}},
+		}}
+		err := secondRunner.RunTransaction(&jujutxn.Transaction{
+			Ops:     ops,
+			Attempt: 0,
+		})
+		c.Assert(err, gc.IsNil)
+	}
+
+	asserted := false
+	hooks := []jujutxn.TestHook{
+		{Asserted: func() {
+			asserted = true
+			setDocName("1", "FooBar")
+		}},
+		{Started: func() { setDocName("1", "Foo") }},
+	}
+	defer txntesting.SetTestHooks(c, s.txnRunner, hooks...).Check()
+	maxAttempt := 0
+	buildTxn := func(attempt int) ([]txn.Op, error) {
+		maxAttempt = attempt
+		ops := []txn.Op{{
+			C:      s.collection.Name,
+			Id:     "1",
+			Assert: bson.D{{"name", "Foo"}},
+			Update: bson.D{{"$set", bson.D{{"name", "Bar"}}}},
+		}}
+		return ops, nil
+	}
+	err := s.txnRunner.Run(buildTxn)
+	c.Assert(err, gc.IsNil)
+	c.Assert(asserted, jc.IsTrue)
+	var found simpleDoc
+	err = s.collection.FindId("1").One(&found)
+	c.Assert(err, gc.IsNil)
+	c.Assert(maxAttempt, gc.Equals, 1)
+	doc := simpleDoc{"1", "Bar"}
+	c.Assert(found, gc.DeepEquals, doc)
+}
+
 func (s *txnSuite) TestAfterHooks(c *gc.C) {
 	changeFuncs := []func(){
 		func() { s.insertDoc(c, "1", "Foo") },