import os
from importlib import import_module

from .testdriver import TestDriver


class Driver(TestDriver):
    drivers = {}

    def __new__(cls, type, *args, **kwargs):
        if type not in cls.drivers:
            try:
                mod = import_module('onitu.drivers.{}.tests.driver'.
                                    format(type))
            except ImportError:
                raise KeyError("No such driver {}".format(repr(type)))
            cls.drivers[type] = mod.Driver
        return cls.drivers[type](*args, **kwargs)


class LocalStorageDriver(TestDriver):
    def __new__(cls, *args, **kwargs):
        return Driver('local_storage', *args, **kwargs)


class TargetDriver(Driver):
    def __new__(cls, *args, **kwargs):
        type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
        return Driver(type, *args, **kwargs)

import os
import pkg_resources

from .testdriver import TestDriver


class Driver(TestDriver):
    drivers = {}

    def __new__(cls, name, *args, **kwargs):
        entry_points = pkg_resources.iter_entry_points('onitu.tests')
        tests_modules = {e.name: e for e in entry_points}

        if name not in tests_modules:
            raise ImportError(
                "Cannot import tests for driver {}".format(name)
            )

        try:
            tests = tests_modules[name].load()
        except ImportError as e:
            raise ImportError(
                "Error importing tests for driver {}: {}".format(name, e)
            )

        try:
            driver = tests.Driver
        except ImportError:
            raise ImportError(
                "Tests for driver {} don't expose a"
                "Driver class".format(name)
            )

        cls.drivers[name] = driver
        return driver(*args, **kwargs)


class LocalStorageDriver(TestDriver):
    def __new__(cls, *args, **kwargs):
        return Driver('local_storage', *args, **kwargs)


class TargetDriver(Driver):
    def __new__(cls, *args, **kwargs):
        type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
        return Driver(type, *args, **kwargs)

Load tests helpers using entry_points
Load tests helpers using entry_points

  import os
- from importlib import import_module
+ import pkg_resources
  
  from .testdriver import TestDriver
  
  
  class Driver(TestDriver):
      drivers = {}
  
-     def __new__(cls, type, *args, **kwargs):
+     def __new__(cls, name, *args, **kwargs):
-         if type not in cls.drivers:
+         entry_points = pkg_resources.iter_entry_points('onitu.tests')
+         tests_modules = {e.name: e for e in entry_points}
+ 
+         if name not in tests_modules:
+             raise ImportError(
+                 "Cannot import tests for driver {}".format(name)
+             )
+ 
-             try:
+         try:
-                 mod = import_module('onitu.drivers.{}.tests.driver'.
-                                     format(type))
+             tests = tests_modules[name].load()
+         except ImportError as e:
+             raise ImportError(
+                 "Error importing tests for driver {}: {}".format(name, e)
+             )
+ 
+         try:
+             driver = tests.Driver
-             except ImportError:
+         except ImportError:
-                 raise KeyError("No such driver {}".format(repr(type)))
+             raise ImportError(
+                 "Tests for driver {} don't expose a"
+                 "Driver class".format(name)
+             )
+ 
-             cls.drivers[type] = mod.Driver
+         cls.drivers[name] = driver
-         return cls.drivers[type](*args, **kwargs)
+         return driver(*args, **kwargs)
  
  
  class LocalStorageDriver(TestDriver):
      def __new__(cls, *args, **kwargs):
          return Driver('local_storage', *args, **kwargs)
  
  
  class TargetDriver(Driver):
      def __new__(cls, *args, **kwargs):
          type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
          return Driver(type, *args, **kwargs)
  
Load tests helpers using entry_points
## Code Before:
import os
from importlib import import_module

from .testdriver import TestDriver


class Driver(TestDriver):
    drivers = {}

    def __new__(cls, type, *args, **kwargs):
        if type not in cls.drivers:
            try:
                mod = import_module('onitu.drivers.{}.tests.driver'.
                                    format(type))
            except ImportError:
                raise KeyError("No such driver {}".format(repr(type)))
            cls.drivers[type] = mod.Driver
        return cls.drivers[type](*args, **kwargs)


class LocalStorageDriver(TestDriver):
    def __new__(cls, *args, **kwargs):
        return Driver('local_storage', *args, **kwargs)


class TargetDriver(Driver):
    def __new__(cls, *args, **kwargs):
        type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
        return Driver(type, *args, **kwargs)

## Instruction:
Load tests helpers using entry_points
## Code After:
import os
import pkg_resources

from .testdriver import TestDriver


class Driver(TestDriver):
    drivers = {}

    def __new__(cls, name, *args, **kwargs):
        entry_points = pkg_resources.iter_entry_points('onitu.tests')
        tests_modules = {e.name: e for e in entry_points}

        if name not in tests_modules:
            raise ImportError(
                "Cannot import tests for driver {}".format(name)
            )

        try:
            tests = tests_modules[name].load()
        except ImportError as e:
            raise ImportError(
                "Error importing tests for driver {}: {}".format(name, e)
            )

        try:
            driver = tests.Driver
        except ImportError:
            raise ImportError(
                "Tests for driver {} don't expose a"
                "Driver class".format(name)
            )

        cls.drivers[name] = driver
        return driver(*args, **kwargs)


class LocalStorageDriver(TestDriver):
    def __new__(cls, *args, **kwargs):
        return Driver('local_storage', *args, **kwargs)


class TargetDriver(Driver):
    def __new__(cls, *args, **kwargs):
        type = os.environ.get('ONITU_TEST_DRIVER', 'local_storage')
        return Driver(type, *args, **kwargs)
