Fixes

calico.py

@@ -96,11 +96,6 @@ def __init__(self, network_config, env):
         The hostname to register endpoints under.
         """
 
-        self.policy_driver = get_policy_driver(self)
-        """
-        Chooses the correct policy driver based on the given configuration
-        """
-
         self.container_engine = get_container_engine(self.k8s_pod_name)
         """
         Chooses the correct container engine based on the given configuration.
@@ -111,7 +106,7 @@ def __init__(self, network_config, env):
         Environment dictionary used when calling the IPAM plugin.
         """
 
-        self.command = env.get(CNI_COMMAND_ENV)
+        self.command = env[CNI_COMMAND_ENV]
         assert self.command in [CNI_CMD_DELETE, CNI_CMD_ADD], \
                 "Invalid CNI command %s" % self.command
         """
@@ -121,22 +116,22 @@ def __init__(self, network_config, env):
           - CNI_CMD_DELETE
         """
 
-        self.container_id = env.get(CNI_CONTAINERID_ENV)
+        self.container_id = env[CNI_CONTAINERID_ENV]
         """
         The container's ID in the containerizer. Required.
         """
 
-        self.cni_netns = env.get(CNI_NETNS_ENV)
+        self.cni_netns = env[CNI_NETNS_ENV]
         """
         Relative path to the network namespace of this container.
         """
 
-        self.interface = env.get(CNI_IFNAME_ENV)
+        self.interface = env[CNI_IFNAME_ENV]
         """
         Name of the interface to create within the container.
         """
 
-        self.cni_path = env.get(CNI_PATH_ENV)
+        self.cni_path = env[CNI_PATH_ENV]
         """
         Path in which to search for CNI plugins.
         """
@@ -163,6 +158,11 @@ def __init__(self, network_config, env):
             # Append any existing args - if they are set.
             self.ipam_env[CNI_ARGS_ENV] += ";%s" % env.get(CNI_ARGS_ENV)
 
+        self.policy_driver = get_policy_driver(self)
+        """
+        Chooses the correct policy driver based on the given configuration
+        """
+
     def execute(self):
         """
         Execute the CNI plugin - uses the given CNI_COMMAND to determine
@@ -349,7 +349,7 @@ def _assign_ips(self, env):
         # Call the IPAM plugin.  Returns the plugin returncode,
         # as well as the CNI result from stdout.
         _log.debug("Assigning IP address")
-        assert env.get(CNI_COMMAND_ENV) == CNI_CMD_ADD
+        assert env[CNI_COMMAND_ENV] == CNI_CMD_ADD
         rc, result = self._call_ipam_plugin(env)
 
         try:
@@ -412,7 +412,7 @@ def _release_ip(self, env):
         :return: None.
         """
         _log.info("Releasing IP address")
-        assert env.get(CNI_COMMAND_ENV) == CNI_CMD_DELETE
+        assert env[CNI_COMMAND_ENV] == CNI_CMD_DELETE
         rc, _ = self._call_ipam_plugin(env)
 
         if rc:
@@ -567,7 +567,7 @@ def _provision_veth(self, endpoint):
             _log.exception("Failed to provision veth interface for endpoint %s",
                            endpoint.name)
             self._remove_workload()
-            self.ipam_env.get[CNI_COMMAND_ENV] = CNI_CMD_DELETE
+            self.ipam_env[CNI_COMMAND_ENV] = CNI_CMD_DELETE
             self._release_ip(self.ipam_env)
             print_cni_error(ERR_CODE_GENERIC, e.message)
             sys.exit(ERR_CODE_GENERIC)

calico_cni/policy_drivers.py

@@ -385,7 +385,7 @@ def get_policy_driver(cni_plugin):
         sys.exit(ERR_CODE_GENERIC)
 
     # Determine which policy driver to use.
-    if cni_plugin.k8s_pod_name:
+    if cni_plugin.running_under_k8s:
         # Running under Kubernetes - decide which Kubernetes driver to use.
         if policy_type == POLICY_MODE_KUBERNETES_ANNOTATIONS or \
            policy_type == POLICY_MODE_KUBERNETES:

tests/unit/test_policy_drivers.py

@@ -20,6 +20,7 @@
 from pycalico.datastore import DatastoreClient
 from pycalico.datastore_datatypes import Endpoint, Rule, Rules
 
+from calico import CniPlugin
 from calico_cni.constants import ERR_CODE_GENERIC
 from calico_cni.policy_drivers import (ApplyProfileError,
                                        get_policy_driver,
@@ -134,7 +135,7 @@ def setUp(self):
                 "policy": {}
         }
         self.driver = KubernetesAnnotationDriver(self.pod_name, self.namespace,
-                self.auth_token, self.api_root)
+                self.auth_token, self.api_root, None, None, None)
         assert_equal(self.driver.profile_name, self.profile_name)
 
         # Mock the DatastoreClient
@@ -351,7 +352,7 @@ def setUp(self):
         self.network_name = "net-name"
         self.namespace = "default"
         self.driver = KubernetesPolicyDriver(self.network_name, self.namespace,
-                                             None, None)
+                                             None, None, None, None, None)
 
         # Mock the DatastoreClient
         self.client = MagicMock(spec=DatastoreClient)
@@ -383,46 +384,53 @@ def test_remove_profile(self):
 class GetPolicyDriverTest(unittest.TestCase): 
 
     def test_get_policy_driver_default_k8s(self):
-        k8s_pod_name = "podname"
-        k8s_namespace = "namespace"
-        config = {"name": "testnetwork"}
-        driver = get_policy_driver(k8s_pod_name, k8s_namespace, config)
+        cni_plugin = Mock(spec=CniPlugin)
+        cni_plugin.network_config = {"name": "testnetwork"}
+        cni_plugin.k8s_pod_name = "podname"
+        cni_plugin.k8s_namespace = "namespace"
+        cni_plugin.running_under_k8s = True
+        driver = get_policy_driver(cni_plugin)
         assert_true(isinstance(driver, KubernetesNoPolicyDriver))
 
     def test_get_policy_driver_k8s_annotations(self):
-        k8s_pod_name = "podname"
-        k8s_namespace = "namespace"
-        config = {"name": "testnetwork"}
-        config["policy"] = {"type": "k8s-annotations"}
-        driver = get_policy_driver(k8s_pod_name, k8s_namespace, config)
+        cni_plugin = Mock(spec=CniPlugin)
+        cni_plugin.network_config = {"name": "testnetwork",
+                                     "policy": {"type": "k8s-annotations"}}
+        cni_plugin.k8s_pod_name = "podname"
+        cni_plugin.k8s_namespace = "namespace"
+        cni_plugin.running_under_k8s = True
+        driver = get_policy_driver(cni_plugin)
         assert_true(isinstance(driver, KubernetesAnnotationDriver))
 
     def test_get_policy_driver_k8s(self):
-        k8s_pod_name = "podname"
-        k8s_namespace = "namespace"
-        config = {"name": "testnetwork"}
-        config["policy"] = {"type": "k8s"}
-        driver = get_policy_driver(k8s_pod_name, k8s_namespace, config)
+        cni_plugin = Mock(spec=CniPlugin)
+        cni_plugin.network_config = {"name": "testnetwork", "policy":{"type": "k8s"}}
+        cni_plugin.k8s_pod_name = "podname"
+        cni_plugin.k8s_namespace = "namespace"
+        cni_plugin.running_under_k8s = True
+        driver = get_policy_driver(cni_plugin)
         assert_true(isinstance(driver, KubernetesPolicyDriver))
 
     def test_get_unknown_policy_driver(self):
         config = {"name": "n", "policy": {"type": "madeup"}}
+        cni_plugin = Mock(spec=CniPlugin)
+        cni_plugin.network_config = config
         with assert_raises(SystemExit) as err:
-            get_policy_driver(None, None, config)
+            get_policy_driver(cni_plugin)
         e = err.exception
         assert_equal(e.code, ERR_CODE_GENERIC)
 
     @patch("calico_cni.policy_drivers.DefaultPolicyDriver", autospec=True)
     def test_get_policy_driver_value_error(self, m_driver):
         # Mock
         m_driver.side_effect = ValueError
-        k8s_pod_name = None
-        k8s_namespace = None 
-        config = {"name": "testnetwork"}
+        cni_plugin = Mock(spec=CniPlugin)
+        cni_plugin.network_config = {"name": "testnetwork"}
+        cni_plugin.running_under_k8s = False
 
         # Call
         with assert_raises(SystemExit) as err:
-            get_policy_driver(k8s_pod_name, k8s_namespace, config)
+            get_policy_driver(cni_plugin)
         e = err.exception
         assert_equal(e.code, ERR_CODE_GENERIC)