Add mesos tests.

tests/unit/test_policy_drivers.py

@@ -26,6 +26,7 @@
                                        get_policy_driver,
                                        PolicyException,
                                        DefaultPolicyDriver,
+                                       MesosPolicyDriver,
                                        KubernetesNoPolicyDriver,
                                        KubernetesAnnotationDriver,
                                        KubernetesPolicyDriver)
@@ -99,6 +100,48 @@ def test_invalid_network_name(self, net_name):
         assert_raises(ValueError, DefaultPolicyDriver, net_name)
 
 
+class MesosPolicyDriverTest(unittest.TestCase):
+    def setUp(self):
+        self.network_name = "test_net_name"
+        self.network_args = {
+          "org.apache.mesos" : {
+            "network_info" : {
+              "name" : "test_net_name",
+              "labels" : {
+                "labels" : [
+                  { "key" : "app", "value" : "myapp" },
+                  { "key" : "env", "value" : "prod" }
+                ]
+              }
+            }
+          }
+        }
+
+        # The init function for Mesos Policy driver will convert these
+        # to a standard dictionary.
+        self.expected_labels = {"app":"myapp", "env":"prod"}
+
+        self.driver = MesosPolicyDriver(self.network_name, self.network_args)
+        assert_equal(self.driver.profile_name, self.network_name)
+        assert_equal(self.driver.labels, self.expected_labels)
+
+        # Mock the datastore client
+        self.client = MagicMock(spec=DatastoreClient)
+        self.driver._client = self.client
+
+    @patch("calico_cni.policy_drivers.DefaultPolicyDriver.apply_profile")
+    def test_apply_profile(self, m_default_apply_profile):
+        endpoint = MagicMock(spec=Endpoint)
+        endpoint.endpoint_id = "12345"
+
+        # Mock default profile call
+        self.driver.apply_profile(endpoint)
+
+        m_default_apply_profile.assert_called_once()
+        self.assertEqual(endpoint.labels, self.expected_labels)
+
+
+
 class KubernetesDefaultPolicyDriverTest(unittest.TestCase):
     """
     Test class for KubernetesDefaultPolicyDriver class.
@@ -122,6 +165,7 @@ def test_generate_rules(self):
                          outbound_rules=[Rule(action="allow")])
         assert_equal(rules, expected)
 
+
 class KubernetesAnnotationDriverTest(unittest.TestCase):
     """
     Test class for KubernetesAnnotationDriver class.
@@ -422,6 +466,7 @@ def test_get_metadata_missing(self):
         # Should be None
         assert_equal(annotations, None)
 
+
 class KubernetesPolicyDriverTest(unittest.TestCase):
     """
     Test class for DefaultDenyInboundDriver class.