adding limits to pycue

pycue/opencue/wrappers/layer.py

@@ -68,6 +68,17 @@ def markdone(self):
         return self.stub.MarkdoneFrames(job_pb2.LayerMarkdoneFramesRequest(layer=self.data),
                                         timeout=Cuebot.Timeout)
 
+    def addLimit(self, limit_id):
+        """Add a limit to the current layer."""
+        return self.stub.AddLimit(job_pb2.LayerAddLimitRequest(layer=self.data, limit_id=limit_id),
+                                  timeout=Cuebot.Timeout)
+
+    def dropLimit(self, limit_id):
+        """Remove a limit on the current layer."""
+        return self.stub.DropLimit(
+            job_pb2.LayerDropLimitRequest(layer=self.data, limit_id=limit_id),
+            timeout=Cuebot.Timeout)
+
     def enableMemoryOptimizer(self, value):
         """Set enableMemoryOptimizer to the value.
         @type value: bool
@@ -252,7 +263,7 @@ def getLimitDetails(self):
         @rtype: list<opencue.wrappers.limit.Limit>
         @return: The list of limits on this layer."""
         return [opencue.wrappers.limit.Limit(limit) for limit in self.stub.GetLimits(
-            job_pb2.LayerGetLimitsRequest(layer=self.data), timeout=Cuebot.Timeout)]
+            job_pb2.LayerGetLimitsRequest(layer=self.data), timeout=Cuebot.Timeout).limits]
 
     def id(self):
         """Returns the uuid of the layer

pycue/opencue/wrappers/limit.py

@@ -35,8 +35,9 @@ def create(self):
     @rtype: opencue.wrappers.limit.Limit
     @return: The newly created Limit
     """
-    return Limit(self.stub.Create(limit_pb2.LimitCreateRequest(name=self.name(),
-                                                               max_value=self.maxValue())))
+    return Limit(self.stub.Create(
+        limit_pb2.LimitCreateRequest(name=self.name(), max_value=self.maxValue()),
+        timeout=Cuebot.Timeout))
 
   def delete(self):
     """Delete the limit record"""
@@ -58,7 +59,7 @@ def get(self, id):
     @rtype: opencue.wrappers.limit.Limit
     @return: The limit found by id.
     """
-    return Limit(self.stub.Get(limit_pb2.LimitGetRequest(id=id), timeout=Cuebot.Timeout))
+    return Limit(self.stub.Get(limit_pb2.LimitGetRequest(id=id), timeout=Cuebot.Timeout).limit)
 
   def rename(self, newName):
     """Rename the current limit to the provided newName.
@@ -96,3 +97,9 @@ def maxValue(self):
       return self.data.max_value
     else:
       return -1
+
+  def currentRunning(self):
+    if hasattr(self.data, 'current_running'):
+      return self.data.current_running
+    else:
+      return -1

pycue/tests/wrappers/layer_test.py

@@ -81,6 +81,34 @@ def testMarkdone(self, getStubMock):
         stubMock.MarkdoneFrames.assert_called_with(
             job_pb2.LayerMarkdoneFramesRequest(layer=layer.data), timeout=mock.ANY)
 
+    def testAddLimit(self, getStubMock):
+        test_limit_id = 'lll-llll-lll'
+        stubMock = mock.Mock()
+        stubMock.AddLimit.return_value = job_pb2.LayerAddLimitResponse()
+        getStubMock.return_value = stubMock
+
+        layer = opencue.wrappers.layer.Layer(
+            job_pb2.Layer(name=TEST_LAYER_NAME))
+        layer.addLimit(test_limit_id)
+
+        stubMock.AddLimit.assert_called_with(
+            job_pb2.LayerAddLimitRequest(layer=layer.data, limit_id=test_limit_id),
+            timeout=mock.ANY)
+
+    def testDropLimit(self, getStubMock):
+        test_limit_id = 'lll-llll-lll'
+        stubMock = mock.Mock()
+        stubMock.DropLimit.return_value = job_pb2.LayerDropLimitResponse()
+        getStubMock.return_value = stubMock
+
+        layer = opencue.wrappers.layer.Layer(
+            job_pb2.Layer(name=TEST_LAYER_NAME))
+        layer.dropLimit(test_limit_id)
+
+        stubMock.DropLimit.assert_called_with(
+            job_pb2.LayerDropLimitRequest(layer=layer.data, limit_id=test_limit_id),
+            timeout=mock.ANY)
+
     def testEnableMemoryOptimizerTrue(self, getStubMock):
         stubMock = mock.Mock()
         stubMock.EnableMemoryOptimizer.return_value = job_pb2.LayerEnableMemoryOptimizerResponse()

pycue/tests/wrappers/limit_test.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+#  Copyright (c) 2018 Sony Pictures Imageworks Inc.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+
+from __future__ import print_function
+from __future__ import division
+from __future__ import absolute_import
+import mock
+import unittest
+
+import opencue
+from opencue.compiled_proto import limit_pb2
+
+
+TEST_LIMIT_ID = 'lll-llll-lll'
+TEST_LIMIT_NAME = 'imalimit'
+TEST_LIMIT_MAX_VALUE = 42
+
+
+@mock.patch('opencue.cuebot.Cuebot.getStub')
+class LimitTests(unittest.TestCase):
+
+    def testCreate(self, getStubMock):
+        stubMock = mock.Mock()
+        stubMock.Create.return_value = limit_pb2.LimitCreateResponse()
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit(
+            limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        limit.create()
+
+        stubMock.Create.assert_called_with(
+            limit_pb2.LimitCreateRequest(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE),
+            timeout=mock.ANY)
+
+    def testDelete(self, getStubMock):
+        stubMock = mock.Mock()
+        stubMock.Delete.return_value = limit_pb2.LimitDeleteResponse()
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit(
+            limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        limit.delete()
+
+        stubMock.Delete.assert_called_with(
+            limit_pb2.LimitDeleteRequest(name=TEST_LIMIT_NAME), timeout=mock.ANY)
+
+    def testFind(self, getStubMock):
+        stubMock = mock.Mock()
+        stubMock.Find.return_value = limit_pb2.LimitFindResponse(
+          limit=limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit()
+        responseLimit = limit.find(TEST_LIMIT_NAME)
+
+        stubMock.Find.assert_called_with(
+            limit_pb2.LimitFindRequest(name=TEST_LIMIT_NAME), timeout=mock.ANY)
+        self.assertEqual(responseLimit.name(), TEST_LIMIT_NAME)
+        self.assertEqual(responseLimit.maxValue(), TEST_LIMIT_MAX_VALUE)
+
+    def testGet(self, getStubMock):
+        stubMock = mock.Mock()
+        stubMock.Get.return_value = limit_pb2.LimitGetResponse(
+            limit=limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit()
+        responseLimit = limit.get(TEST_LIMIT_ID)
+
+        stubMock.Get.assert_called_with(
+            limit_pb2.LimitGetRequest(id=TEST_LIMIT_ID), timeout=mock.ANY)
+        self.assertEqual(responseLimit.name(), TEST_LIMIT_NAME)
+        self.assertEqual(responseLimit.maxValue(), TEST_LIMIT_MAX_VALUE)
+
+    def testRename(self, getStubMock):
+        test_new_name = 'new_name'
+        stubMock = mock.Mock()
+        stubMock.Rename.return_value = limit_pb2.LimitRenameResponse()
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit(
+            limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        limit.rename(test_new_name)
+
+        stubMock.Rename.assert_called_with(
+            limit_pb2.LimitRenameRequest(old_name=TEST_LIMIT_NAME, new_name=test_new_name),
+            timeout=mock.ANY)
+
+    def testSetMaxValue(self, getStubMock):
+        max_value = 16
+        stubMock = mock.Mock()
+        stubMock.SetMaxValue.return_value = limit_pb2.LimitSetMaxValueResponse()
+        getStubMock.return_value = stubMock
+
+        limit = opencue.wrappers.limit.Limit(
+            limit_pb2.Limit(name=TEST_LIMIT_NAME, max_value=TEST_LIMIT_MAX_VALUE))
+        limit.setMaxValue(max_value)
+
+        stubMock.SetMaxValue.assert_called_with(
+            limit_pb2.LimitSetMaxValueRequest(name=TEST_LIMIT_NAME, max_value=max_value),
+            timeout=mock.ANY)
+
+
+if __name__ == '__main__':
+  unittest.main()