test the output_layer argument passed to the callable model wrapper

tests_tf/test_attacks.py

@@ -13,6 +13,8 @@
 from cleverhans.attacks import VirtualAdversarialMethod
 from cleverhans.attacks import SaliencyMapMethod
 
+from cleverhans.model import CallableModelWrapper
+
 
 class TestAttackClassInitArguments(unittest.TestCase):
     def test_model(self):
@@ -291,6 +293,19 @@ def test_generate_np_targeted_gives_adversarial_example(self):
         worked = np.mean(np.argmax(feed_labs, axis=1) == new_labs)
         self.assertTrue(worked > .9)
 
+class TestCallableModelWrapperInitArguments(unittest.TestCase):
+    def test_output_layer(self):
+        def model():
+            return True
+
+        # The following two calls should not raise Exceptions
+        wrap = CallableModelWrapper(model, 'probs')
+        wrap = CallableModelWrapper(model, 'logits')
+
+        with self.assertRaises(Exception) as context:
+            wrap = CallableModelWrapper(model, 'dummy')
+        self.assertTrue(context.exception)
+
 
 if __name__ == '__main__':
     unittest.main()