diff --git a/src/main/java/cz/muni/ics/kypo/training/adaptive/domain/training/TrainingRun.java b/src/main/java/cz/muni/ics/kypo/training/adaptive/domain/training/TrainingRun.java
index 85444b6a83cd63ab474e918046ff90016265e4a0..2ba5fa31f4220b64b9e92c0bca77acbdc2333376 100644
--- a/src/main/java/cz/muni/ics/kypo/training/adaptive/domain/training/TrainingRun.java
+++ b/src/main/java/cz/muni/ics/kypo/training/adaptive/domain/training/TrainingRun.java
@@ -131,7 +131,7 @@ public class TrainingRun implements Serializable {
     @Column(name = "sandbox_instance_ref_id", length = 36)
     private String sandboxInstanceRefId;
     @Column(name = "sandbox_instance_allocation_id")
-    private Long sandboxInstanceAllocationId;
+    private Integer sandboxInstanceAllocationId;
     @ManyToOne(fetch = FetchType.LAZY)
     @JoinColumn(name = "user_id", nullable = false)
     private User participantRef;
@@ -287,7 +287,7 @@ public class TrainingRun implements Serializable {
      *
      * @return the sandbox instance allocation id
      */
-    public Long getSandboxInstanceAllocationId() {
+    public Integer getSandboxInstanceAllocationId() {
         return sandboxInstanceAllocationId;
     }
 
@@ -295,7 +295,7 @@ public class TrainingRun implements Serializable {
      * Sets sandbox instance allocation id associated with Training run
      * @param sandboxInstanceAllocationId the sandbox instance allocation id
      */
-    public void setSandboxInstanceAllocationId(Long sandboxInstanceAllocationId) {
+    public void setSandboxInstanceAllocationId(Integer sandboxInstanceAllocationId) {
         this.sandboxInstanceAllocationId = sandboxInstanceAllocationId;
     }
 
diff --git a/src/main/java/cz/muni/ics/kypo/training/adaptive/dto/trainingrun/TrainingRunDTO.java b/src/main/java/cz/muni/ics/kypo/training/adaptive/dto/trainingrun/TrainingRunDTO.java
index ad72ceb37664f012e177c099c32c57f5256cdd58..66be6ef71e5f87d0aad63e31b9f4895cf8cbe176 100644
--- a/src/main/java/cz/muni/ics/kypo/training/adaptive/dto/trainingrun/TrainingRunDTO.java
+++ b/src/main/java/cz/muni/ics/kypo/training/adaptive/dto/trainingrun/TrainingRunDTO.java
@@ -32,7 +32,7 @@ public class TrainingRunDTO {
     @ApiModelProperty(value = "Reference to the received sandbox.")
     private String sandboxInstanceRefId;
     @ApiModelProperty(value = "Allocation id to the received sandbox.")
-    private Long sandboxInstanceAllocationId;
+    private Integer sandboxInstanceAllocationId;
     @ApiModelProperty(value = "Reference to participant of training run.")
     private UserRefDTO participantRef;
     @ApiModelProperty(value = "Boolean to check whether event logging works.", example = "true")
@@ -135,7 +135,7 @@ public class TrainingRunDTO {
      *
      * @return the sandbox instance allocation id
      */
-    public Long getSandboxInstanceAllocationId() {
+    public Integer getSandboxInstanceAllocationId() {
         return sandboxInstanceAllocationId;
     }
 
@@ -144,7 +144,7 @@ public class TrainingRunDTO {
      *
      * @param sandboxInstanceAllocationId the sandbox instance allocation id
      */
-    public void setSandboxInstanceAllocationId(Long sandboxInstanceAllocationId) {
+    public void setSandboxInstanceAllocationId(Integer sandboxInstanceAllocationId) {
         this.sandboxInstanceAllocationId = sandboxInstanceAllocationId;
     }
 
diff --git a/src/main/java/cz/muni/ics/kypo/training/adaptive/service/training/TrainingRunService.java b/src/main/java/cz/muni/ics/kypo/training/adaptive/service/training/TrainingRunService.java
index 7bfb18a6a7f491e0fe694474ec663fb77b3f58b8..b99c404ec4ce84cc1dcff4bc6840b6535d7be716 100644
--- a/src/main/java/cz/muni/ics/kypo/training/adaptive/service/training/TrainingRunService.java
+++ b/src/main/java/cz/muni/ics/kypo/training/adaptive/service/training/TrainingRunService.java
@@ -605,8 +605,9 @@ public class TrainingRunService {
      */
     public TrainingRun assignSandbox(TrainingRun trainingRun, long poolId) {
         String sandboxInstanceRef = this.sandboxServiceApi.getAndLockSandboxForTrainingRun(poolId);
-        Long sandboxInstanceAllocationId = this.sandboxServiceApi.getAndLockSandbox(poolId).getAllocationUnitId().longValue();
+        Integer sandboxInstanceAllocationId = this.sandboxServiceApi.getAndLockSandbox(poolId).getAllocationUnitId();
         trainingRun.setSandboxInstanceRefId(sandboxInstanceRef);
+        trainingRun.setSandboxInstanceAllocationId(sandboxInstanceAllocationId);
         return trainingRunRepository.save(trainingRun);
     }
 
diff --git a/src/test/java/cz/muni/ics/kypo/training/adaptive/util/TestDataFactory.java b/src/test/java/cz/muni/ics/kypo/training/adaptive/util/TestDataFactory.java
index 7d4c7365e1c55f3b278fe42122f53343448f732c..eed6cc87eabaf1c8cb2023e76dd32512669bdc0a 100644
--- a/src/test/java/cz/muni/ics/kypo/training/adaptive/util/TestDataFactory.java
+++ b/src/test/java/cz/muni/ics/kypo/training/adaptive/util/TestDataFactory.java
@@ -177,16 +177,16 @@ public class TestDataFactory {
             LocalDateTime.now(Clock.systemUTC()).minusHours(10), "Archived instance", "archived-6666");
 
     private TrainingRun runningRun = generateTrainingRun(LocalDateTime.now(Clock.systemUTC()).minusHours(2), LocalDateTime.now(Clock.systemUTC()).plusHours(2),
-            TRState.RUNNING, 2, true, "1L", 1L, true,
+            TRState.RUNNING, 2, true, "1L", 1, true,
             "20L");
     private TrainingRun finishedRun = generateTrainingRun(LocalDateTime.now(Clock.systemUTC()).minusHours(10), LocalDateTime.now(Clock.systemUTC()).minusHours(5),
-            TRState.FINISHED, 4, false, "3L", 3l, true, "30L");
+            TRState.FINISHED, 4, false, "3L", 3, true, "30L");
     private TrainingRun archivedRun = generateTrainingRun(LocalDateTime.now(Clock.systemUTC()).minusHours(20), LocalDateTime.now(Clock.systemUTC()).minusHours(10),
-            TRState.ARCHIVED, 0, false, "5L", 5l, false, "60L");
+            TRState.ARCHIVED, 0, false, "5L", 5, false, "60L");
     private TrainingRunByIdDTO trainingRunByIdDTO = generateTrainingRunByIdDTO(LocalDateTime.now(Clock.systemUTC()).minusHours(2), LocalDateTime.now(Clock.systemUTC()).plusHours(2),
              TRState.RUNNING, "5L");
     private TrainingRunDTO trainingRunDTO = generateTrainingRunDTO(LocalDateTime.now(Clock.systemUTC()).minusHours(9), LocalDateTime.now(Clock.systemUTC()).minusHours(5),
-             TRState.FINISHED, "7L", 7L);
+             TRState.FINISHED, "7L", 7);
 //    private AccessedTrainingRunDTO accessedTrainingRunDTO = generateAccessedTrainingRunDTO("Accessed run", LocalDateTime.now(Clock.systemUTC()).minusHours(8), LocalDateTime.now(Clock.systemUTC()).minusHours(4), 5,
 //            6, Actions.RESUME);
 
@@ -694,7 +694,7 @@ public class TestDataFactory {
     }
 
     private TrainingRun generateTrainingRun(LocalDateTime startTime, LocalDateTime endTime, TRState state,
-                                            int incorrectAnswerCount, boolean solutionTaken, String SBIRefId, Long SBIAllocId, boolean phaseAnswered, String previousSBIRefId){
+                                            int incorrectAnswerCount, boolean solutionTaken, String SBIRefId, Integer SBIAllocId, boolean phaseAnswered, String previousSBIRefId){
         TrainingRun newTrainingRun = new TrainingRun();
         newTrainingRun.setStartTime(startTime);
         newTrainingRun.setEndTime(endTime);
@@ -805,7 +805,7 @@ public class TestDataFactory {
         return trainingRunByIdDTO;
     }
 
-    private TrainingRunDTO generateTrainingRunDTO(LocalDateTime start, LocalDateTime end, TRState state, String SBIRefId, Long SBIAllocId){
+    private TrainingRunDTO generateTrainingRunDTO(LocalDateTime start, LocalDateTime end, TRState state, String SBIRefId, Integer SBIAllocId){
         TrainingRunDTO trainingRunDTO = new TrainingRunDTO();
         trainingRunDTO.setStartTime(start);
         trainingRunDTO.setEndTime(end);