selftest.target: Add tests for setup.

2 files changed, 14 insertions, 9 deletions

diff --git a/selftest/target/__init__.py b/selftest/target/__init__.py
index 03a440dd27..ac2cc19a48 100644
--- a/selftest/target/__init__.py
+++ b/selftest/target/__init__.py
@@ -91,13 +91,15 @@ class EnvironmentManager(object):
 
     def get_running_env(self, name):
         envname = name.split(":")[0]
-        return self.running_envs[envname]
+        if envname == "none":
+            return NoneEnvironment()
+        return self.running_envs.get(envname)
 
     def getlog_env(self, envname):
         env = self.get_running_env(envname)
         return env.get_log()
 
-    def check_env(envname):
+    def check_env(self, envname):
         """Check if an environment is still up.
 
         :param envname: Environment to check
@@ -115,13 +117,11 @@ class EnvironmentManager(object):
         del self.running_envs[envname]
 
     def teardown_all(self):
+        """Teardown all environments."""
         for env in self.running_envs.iterkeys():
             self.teardown_env(env)
 
     def setup_env(self, envname, prefix):
-        if envname == "none":
-            return NoneEnvironment()
-
         running_env = self.get_running_env(envname)
         if running_env is not None:
             if not running_env.check():
diff --git a/selftest/tests/test_target.py b/selftest/tests/test_target.py
index 9b3d37b270..343de930cb 100644
--- a/selftest/tests/test_target.py
+++ b/selftest/tests/test_target.py
@@ -26,7 +26,6 @@ from selftest.target import (
     Target,
     )
 
-import os
 import unittest
 
 
@@ -47,8 +46,8 @@ class DummyEnvironment(Environment):
 
 class DummyTarget(Target):
 
-    def get_target(self, name, prefix):
-        return DummyTarget(name, prefix)
+    def setup_env(self, name, prefix):
+        return DummyEnvironment(name, prefix)
 
 
 class EnvironmentManagerTests(unittest.TestCase):
@@ -57,4 +56,10 @@ class EnvironmentManagerTests(unittest.TestCase):
         self.mgr = EnvironmentManager(DummyTarget())
 
     def test_none(self):
-        self.assertIs(NoneEnvironment, type(self.mgr.setup_env("none", "prefix")))
+        self.assertIs(
+            NoneEnvironment, type(self.mgr.setup_env("none", "prefix")))
+
+    def test_setup(self):
+        env = self.mgr.setup_env("something", "prefix")
+        self.assertEquals(env.name, "something")
+        self.assertEquals(env.prefix, "prefix")