diff --git a/tests/TestCase.php b/tests/TestCase.php
index 3b9ba78..0f0bd48 100644
--- a/tests/TestCase.php
+++ b/tests/TestCase.php
@@ -4,7 +4,7 @@ namespace yiiunit\extensions\swiftmailer;
 
 use yii\helpers\ArrayHelper;
 
-abstract class TestCase extends \PHPUnit_Framework_TestCase
+abstract class TestCase extends \PHPUnit\Framework\TestCase
 {
     /**
      * Clean up after test.
diff --git a/tests/compatibility.php b/tests/compatibility.php
index 9141b4f..152645a 100644
--- a/tests/compatibility.php
+++ b/tests/compatibility.php
@@ -1,12 +1,16 @@
 <?php
 /*
- * Ensures compatibility with PHPUnit 6.x
+ * Ensures compatibility with PHPUnit < 6.x
  */
 
-if (!class_exists('PHPUnit_Framework_Constraint') && class_exists('PHPUnit\Framework\Constraint\Constraint')) {
-    abstract class PHPUnit_Framework_Constraint extends \PHPUnit\Framework\Constraint\Constraint {}
+if (!class_exists('PHPUnit\Framework\Constraint\Constraint') && class_exists('PHPUnit_Framework_Constraint')) {
+    namespace PHPUnit\Framework\Constraint {
+        abstract class Constraint extends \PHPUnit_Framework_Constraint {}
+    }
 }
 
-if (!class_exists('PHPUnit_Framework_TestCase') && class_exists('PHPUnit\Framework\TestCase')) {
-    abstract class PHPUnit_Framework_TestCase extends \PHPUnit\Framework\TestCase {}
+if (!class_exists('PHPUnit\Framework\TestCase') && class_exists('PHPUnit_Framework_TestCase')) {
+    namespace PHPUnit\Framework {
+        abstract class TestCase extends \PHPUnit_Framework_TestCase {}
+    }
 }