Cleaning up EF 6 tests with Transaction Rollbacks

You’ve set up Entity Framework, and you’ve even written integration tests against your code. Your tests create all kinds of sample data in your test database (you aren’t running integration tests against your production database, are you?), and now you need to make sure that the sample data gets cleaned up so the database looks like it did before your test. You could try to do some code to cleanup your data after your test runs, but if your test fails, it can be hard to know exactly how to clean up your data because you can’t guarantee that the data is in a particular state.

There is a better way. I was chatting with @jeremy6d this morning, and he suggested that I use transactions that can be rolled back at the end of a test. The following is how to do so when using Entity Framework 6.

Disable Retry

The first thing to be aware of is that if you want to use a transaction with Entity Framework, you can’t use the Retrying Execution Strategy. If you’re using that strategy in your project, you’ll need to turn it off for your tests. To do this, first look at your DBContext and see if you have a retry strategy set up. It will likely be an attribute on the class, like this:

1
2
[DbConfigurationType(typeof(RetryDbConfiguration))]
public partial class EntitiesV3 : DbContext

Our RetryDbConfiguration looked like this:

1
2
3
4
5
6
7
8
public class RetryDbConfiguration : DbConfiguration
{
public RetryDbConfiguration()
{

SetExecutionStrategy("System.Data.SqlClient", () =>
new SqlAzureExecutionStrategy(5, TimeSpan.FromSeconds(15)));
}
}

I then added in some properties to let me disable the Retry Execution Strategy.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
public class RetryDbConfiguration : DbConfiguration
{
public RetryDbConfiguration()
{

this.SetExecutionStrategy("System.Data.SqlClient", () =>
SuspendExecutionStrategy
? (IDbExecutionStrategy)new DefaultExecutionStrategy()
: new SqlAzureExecutionStrategy(5, TimeSpan.FromSeconds(15)));
}

public static bool SuspendExecutionStrategy
{
get
{
return (bool?)CallContext.LogicalGetData
("SuspendExecutionStrategy") ?? false;
}
set
{
CallContext.LogicalSetData("SuspendExecutionStrategy", value);
}
}
}

Implement Transactions

Next, create a reusable base class for your test classes.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
[TestClass]
public class TransactionTest
{
protected EntitiesV3 context;
protected DbContextTransaction transaction;

[AssemblyInitialize]
public static void AssemblyStart(TestContext testContext)
{

RetryDbConfiguration.SuspendExecutionStrategy = true;
}

[TestInitialize]
public void TransactionTestStart()
{

context = new EntitiesV3();
transaction = context.Database.BeginTransaction();
}

[TestCleanup]
public void TransactionTestEnd()
{

transaction.Rollback();
transaction.Dispose();
context.Dispose();
}

[AssemblyCleanup]
public static void AssemblyEnd()
{

RetryDbConfiguration.SuspendExecutionStrategy = false;
}
}

This will create a transaction before each tests starts and rollback the transaction after each test ends.

Update 2014-08-19: Please note that the name of the [TestInitialize] and [TestCleanup] methods must be different from the name of the [TestInitialize] and [TestCleanup] methods in any derived classes, otherwise the ones in the base class will be ignored. I have updated the names of my methods in this example to TransactionTestStart and TransactionTestEnd to try to avoid any such naming conflicts.

Update 2014-08-21: I moved the context creation and disposal inside the [TestInitialize] and [TestCleanup], because I found that the transaction rollback does not remove entities that were cached by the context, which causes problems for subsequent tests. Creating a new context for each test solves this problem.

Write Tests

In your tests, you can then create objects, call context.SaveChanges() in order to generate and get access to keys, and do all kinds of other things that require reading from and writing to the database. When your test is done, this code will then roll back the transaction and the database will look just like it did before your test ran. Here is what a test might look like:

1
2
3
4
5
6
7
8
9
10
11
[TestClass]
public class ItemTests: TransactionTest
{
[TestMethod]public void GetItemShouldDoSomething()
{

context.Items.Add(new Item() {Name = "Thing"});
context.SaveChanges();
var item = context.Items.FirstOrDefault(i => i.Name == "Thing");
item.Should().NotBeNull();
item.Id.Should().Not().Be(0);
}

Notice that we are calling SaveChanges, then later we are retrieving an item from the database and checking that it’s ID property was generated by the database. All of this will be rolled back.

If you’re curious, I’m using Fluent Assertions to make the last part of the test a little easier to read.

References

Transactions in Entity Framework 6: http://msdn.microsoft.com/en-us/data/dn456843.aspx
Transactions in older versions of Entity Framework: http://mrpanot.wordpress.com/2010/11/25/use-transactions-and-rollback-for-unit-testing-against-entity-framework/
Using a base Test Class: http://blogs.msdn.com/b/densto/archive/2008/05/16/using-a-base-class-for-your-unit-test-classes.aspx
Limitations when using a base Test Class: https://connect.microsoft.com/VisualStudio/feedback/details/687481/classinitialize-attribute-not-used-in-a-based-class
Suspending EF Execution Strategy: http://msdn.microsoft.com/en-US/data/dn307226