diff --git a/src/SplitTestsByGroups.php b/src/SplitTestsByGroups.php
index 760d3c3..75e33c9 100644
--- a/src/SplitTestsByGroups.php
+++ b/src/SplitTestsByGroups.php
@@ -3,7 +3,8 @@
 
 use Codeception\Test\Descriptor as TestDescriptor;
 use Codeception\Test\Loader as TestLoader;
-use \PHPUnit_Framework_TestSuite_DataProvider as DataProviderTestSuite;
+use \PHPUnit\Framework\DataProviderTestSuite as DataProviderTestSuite;
+use \PHPUnit_Framework_TestSuite_DataProvider as DataProvider;
 use Robo\Contract\TaskInterface;
 use Robo\Exception\TaskException;
 use Robo\Task\BaseTask;
@@ -94,7 +95,7 @@ public function run()
         $this->printTaskInfo('Processing ' . count($tests) . ' tests');
         // splitting tests by groups
         foreach ($tests as $test) {
-            if ($test instanceof DataProviderTestSuite) {
+            if ($test instanceof DataProvider || $test instanceof DataProviderTestSuite) {
                 $test = current($test->tests());
             }
             $groups[($i % $this->numGroups) + 1][] = TestDescriptor::getTestFullName($test);