Add GetDefault and SetDefault to AllcationInterface (#939)

AcademySoftwareFoundation · Apr 25, 2021 · 847ee50 · 847ee50
1 parent bb2c252
commit 847ee50
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 2 deletions.
diff --git a/VERSION.in b/VERSION.in
@@ -1 +1 @@
-0.11
+0.12
diff --git a/cuebot/src/main/java/com/imageworks/spcue/servant/ManageAllocation.java b/cuebot/src/main/java/com/imageworks/spcue/servant/ManageAllocation.java
@@ -24,6 +24,8 @@
 import com.imageworks.spcue.grpc.facility.AllocFindResponse;
 import com.imageworks.spcue.grpc.facility.AllocGetAllRequest;
 import com.imageworks.spcue.grpc.facility.AllocGetAllResponse;
+import com.imageworks.spcue.grpc.facility.AllocGetDefaultRequest;
+import com.imageworks.spcue.grpc.facility.AllocGetDefaultResponse;
 import com.imageworks.spcue.grpc.facility.AllocGetHostsRequest;
 import com.imageworks.spcue.grpc.facility.AllocGetHostsResponse;
 import com.imageworks.spcue.grpc.facility.AllocGetRequest;
@@ -34,6 +36,8 @@
 import com.imageworks.spcue.grpc.facility.AllocReparentHostsResponse;
 import com.imageworks.spcue.grpc.facility.AllocSetBillableRequest;
 import com.imageworks.spcue.grpc.facility.AllocSetBillableResponse;
+import com.imageworks.spcue.grpc.facility.AllocSetDefaultRequest;
+import com.imageworks.spcue.grpc.facility.AllocSetDefaultResponse;
 import com.imageworks.spcue.grpc.facility.AllocSetNameRequest;
 import com.imageworks.spcue.grpc.facility.AllocSetNameResponse;
 import com.imageworks.spcue.grpc.facility.AllocSetTagRequest;
@@ -221,6 +225,28 @@ public void setTag(
  responseObserver.onCompleted();
  }
 
+ @Override
+ public void getDefault(
+ AllocGetDefaultRequest request,
+ StreamObserver<AllocGetDefaultResponse> responseObserver) {
+ AllocationEntity alloc = adminManager.getDefaultAllocation();
+ responseObserver.onNext(AllocGetDefaultResponse.newBuilder()
+ .setAllocation(whiteboard.getAllocation(alloc.id))
+ .build());
+ responseObserver.onCompleted();
+ }
+
+ @Override
+ public void setDefault(
+ AllocSetDefaultRequest request,
+ StreamObserver<AllocSetDefaultResponse> responseObserver) {
+ AllocationEntity alloc = adminManager.findAllocationDetail(
+ request.getAllocation().getFacility(), request.getAllocation().getName());
+ adminManager.setDefaultAllocation(alloc);
+ responseObserver.onNext(AllocSetDefaultResponse.newBuilder().build());
+ responseObserver.onCompleted();
+ }
+
  public AdminManager getAdminManager() {
  return adminManager;
  }

diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/AdminManager.java b/cuebot/src/main/java/com/imageworks/spcue/service/AdminManager.java
@@ -56,6 +56,8 @@ public interface AdminManager {
  void deleteAllocation(AllocationInterface alloc);
  void setAllocationName(AllocationInterface a, String name);
  void setAllocationTag(AllocationInterface a, String tag);
+ AllocationEntity getDefaultAllocation();
+ void setDefaultAllocation(AllocationInterface a);
  AllocationEntity findAllocationDetail(String facility, String name);
  AllocationEntity getAllocationDetail(String id);
  void setAllocationBillable(AllocationInterface alloc, boolean value);

diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/AdminManagerService.java b/cuebot/src/main/java/com/imageworks/spcue/service/AdminManagerService.java
@@ -106,6 +106,17 @@ public void setAllocationTag(AllocationInterface a, String tag) {
  allocationDao.updateAllocationTag(a, tag);
  }
 
+ @Override
+ @Transactional(propagation = Propagation.REQUIRED, readOnly=true)
+ public AllocationEntity getDefaultAllocation() {
+ return allocationDao.getDefaultAllocationEntity();
+ }
+
+ @Override
+ public void setDefaultAllocation(AllocationInterface a) {
+ allocationDao.setDefaultAllocation(a);
+ }
+
  @Override
  @Transactional(propagation = Propagation.REQUIRED, readOnly=true)
  public ShowEntity findShowEntity(String name) {

diff --git a/cuebot/src/test/java/com/imageworks/spcue/test/service/AdminManagerTests.java b/cuebot/src/test/java/com/imageworks/spcue/test/service/AdminManagerTests.java
@@ -65,6 +65,34 @@ public void createAllocation() {
  adminManager.createAllocation(facilityDao.getDefaultFacility(), a);
  }
 
+ @Test
+ @Transactional
+ @Rollback(true)
+ public void deleteAllocation() {
+ AllocationEntity a = new AllocationEntity();
+ a.name = facilityDao.getDefaultFacility().getName() + "." + TEST_ALLOC_NAME;
+ a.tag = "general";
+ adminManager.createAllocation(facilityDao.getDefaultFacility(), a);
+ adminManager.deleteAllocation(a);
+ }
+
+ @Test
+ @Transactional
+ @Rollback(true)
+ public void setDefaultAllocation() {
+ AllocationEntity a = adminManager.getDefaultAllocation();
+ assertEquals(a.name, facilityDao.getDefaultFacility().getName() + ".unassigned");
+
+ a = new AllocationEntity();
+ a.name = TEST_ALLOC_NAME;
+ a.tag = "general";
+ adminManager.createAllocation(facilityDao.getDefaultFacility(), a);
+ adminManager.setDefaultAllocation(a);
+
+ a = adminManager.getDefaultAllocation();
+ assertEquals(a.name, facilityDao.getDefaultFacility().getName() + "." + TEST_ALLOC_NAME);
+ }
+
  @Test
  @Transactional
  @Rollback(true)

diff --git a/proto/facility.proto b/proto/facility.proto
@@ -62,6 +62,12 @@ service AllocationInterface {
 
  // Set the allocation tag. Setting this will re-tag all the hosts in this allocation.
  rpc SetTag(AllocSetTagRequest) returns (AllocSetTagResponse);
+
+ // Return the default allocation.
+ rpc GetDefault(AllocGetDefaultRequest) returns (AllocGetDefaultResponse);
+
+ // Set the default allocation.
+ rpc SetDefault(AllocSetDefaultRequest) returns (AllocSetDefaultResponse);
 }
 
 
@@ -245,4 +251,18 @@ message AllocSetTagRequest {
  string tag = 2;
 }
 
-message AllocSetTagResponse {} //
+message AllocSetTagResponse {} // Empty
+
+// GetDefault
+message AllocGetDefaultRequest {} // Empty
+
+message AllocGetDefaultResponse {
+ Allocation allocation = 1;
+}
+
+// SetDefault
+message AllocSetDefaultRequest {
+ Allocation allocation = 1;
+}
+
+message AllocSetDefaultResponse {} // Empty
diff --git a/pycue/opencue/api.py b/pycue/opencue/api.py
@@ -644,6 +644,28 @@ def deleteAllocation(alloc):
  facility_pb2.AllocDeleteRequest(allocation=alloc), timeout=Cuebot.Timeout)
 
 
+@util.grpcExceptionParser
+def getDefaultAllocation():
+ """Get the default allocation.
+
+ :rtype: Allocation
+ :return: an Allocation object"""
+ return Allocation(Cuebot.getStub('allocation').GetDefault(
+ facility_pb2.AllocGetDefaultRequest(), timeout=Cuebot.Timeout).allocation)
+
+
+@util.grpcExceptionParser
+def setDefaultAllocation(alloc):
+ """Set the default allocation.
+
+ :type alloc: facility_pb2.Allocation
+ :param alloc: allocation to set default
+ :rtype: facility_pb2.AllocSetDefaultResponse
+ :return: empty response"""
+ return Cuebot.getStub('allocation').SetDefault(
+ facility_pb2.AllocSetDefaultRequest(allocation=alloc), timeout=Cuebot.Timeout)
+
+
 @util.grpcExceptionParser
 def allocSetBillable(alloc, is_billable):
  """Sets an allocation billable or not.

diff --git a/pycue/tests/api_test.py b/pycue/tests/api_test.py
@@ -719,6 +719,32 @@ def testDeleteAlloc(self, getStubMock):
  stubMock.Delete.assert_called_with(
  facility_pb2.AllocDeleteRequest(allocation=allocToDelete), timeout=mock.ANY)
 
+ @mock.patch('opencue.cuebot.Cuebot.getStub')
+ def testGetDefaultAlloc(self, getStubMock):
+ arbitraryId = '00000000-0000-0000-0000-012345678980'
+ stubMock = mock.Mock()
+ stubMock.GetDefault.return_value = facility_pb2.AllocGetDefaultResponse(
+ allocation=facility_pb2.Allocation(id=arbitraryId))
+ getStubMock.return_value = stubMock
+
+ alloc = opencue.api.getDefaultAllocation()
+
+ stubMock.GetDefault.assert_called_with(
+ facility_pb2.AllocGetDefaultRequest(), timeout=mock.ANY)
+ self.assertEqual(arbitraryId, alloc.id())
+
+ @mock.patch('opencue.cuebot.Cuebot.getStub')
+ def testSetDefaultAlloc(self, getStubMock):
+ alloc = facility_pb2.Allocation(name=TEST_ALLOC_NAME)
+ stubMock = mock.Mock()
+ stubMock.SetDefault.return_value = facility_pb2.AllocSetDefaultResponse()
+ getStubMock.return_value = stubMock
+
+ opencue.api.setDefaultAllocation(alloc)
+
+ stubMock.SetDefault.assert_called_with(
+ facility_pb2.AllocSetDefaultRequest(allocation=alloc), timeout=mock.ANY)
+
  @mock.patch('opencue.cuebot.Cuebot.getStub')
  def testAllocSetBillable(self, getStubMock):
  alloc = facility_pb2.Allocation(name=TEST_ALLOC_NAME)