5 files changed, 22 insertions, 4 deletions
diff --git a/tests/fdx_test.py b/tests/fdx_test.py
index 8034f71..b6e2bc9 100644
--- a/tests/fdx_test.py
+++ b/tests/fdx_test.py
@@ -2,7 +2,7 @@
 # Licensed under the MIT license:
 # http://www.opensource.org/licenses/mit-license.php
 
-from unittest import TestCase
+from testcompat import TestCase
 from StringIO import StringIO
 
 from screenplain.export.fdx import write_text
diff --git a/tests/files_test.py b/tests/files_test.py
index dc799bc..b5a390c 100644
--- a/tests/files_test.py
+++ b/tests/files_test.py
@@ -4,7 +4,7 @@
 
 from __future__ import with_statement
 
-from unittest import TestCase
+from testcompat import TestCase
 import tempfile
 import os
 import os.path
diff --git a/tests/fountain_test.py b/tests/fountain_test.py
index 5095e4e..e265da9 100644
--- a/tests/fountain_test.py
+++ b/tests/fountain_test.py
@@ -2,7 +2,8 @@
 # Licensed under the MIT license:
 # http://www.opensource.org/licenses/mit-license.php
 
-from unittest import TestCase
+from testcompat import TestCase
+
 from screenplain.parsers import fountain
 from screenplain.types import (
     Slug, Action, Dialog, DualDialog, Transition, Section, PageBreak
diff --git a/tests/richstring_test.py b/tests/richstring_test.py
index b9ca3d2..f2bed5e 100644
--- a/tests/richstring_test.py
+++ b/tests/richstring_test.py
@@ -2,7 +2,7 @@
 # Licensed under the MIT license:
 # http://www.opensource.org/licenses/mit-license.php
 
-from unittest import TestCase
+from testcompat import TestCase
 from screenplain.richstring import (
     RichString, Segment,
     Bold, Italic,
diff --git a/tests/testcompat.py b/tests/testcompat.py
new file mode 100644
index 0000000..e0d0cba
--- /dev/null
+++ b/tests/testcompat.py
@@ -0,0 +1,17 @@
+# Copyright (c) 2011 Martin Vilcans
+# Licensed under the MIT license:
+# http://www.opensource.org/licenses/mit-license.php
+
+"""Provides compatibility between Python 2 and 3.
+
+In Python 2 we use the unittest2 module.
+The functionality of that module is already in Python 3,
+so we don't depend on it.
+This module exports the TestCase class from whatever unittest library we have.
+
+"""
+
+try:
+    from unittest2 import TestCase
+except ImportError:
+    from unittest import TestCase