From 1ec151a138cd594e2048835059276da16719cd6d Mon Sep 17 00:00:00 2001
From: Siratee K <siratee.k@lmwn.com>
Date: Mon, 25 Mar 2024 14:48:04 +0700
Subject: [PATCH] fix: Update Example Repo

---
 main.go | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index bbe1190..2930d10 100644
--- a/main.go
+++ b/main.go
@@ -10,16 +10,36 @@ import (
 	"time"
 )
 
+type TaskAttribute struct {
+	SleepTime            *time.Duration `json:"sleepTime,omitempty"`
+	MemoryAllocationSize *int           `json:"memoryAllocationSize,omitempty"`
+}
+
 func main() {
 	containerlib.Run(func(ctx containerlibcontext.Context, task model.Task) error {
-		var unmarshalledData map[string]interface{}
+		var unmarshalledData TaskAttribute
 		err := json.Unmarshal(task.Attributes, &unmarshalledData)
 		if err != nil {
 			logrus.Error(err)
 			return err
 		}
 
-		time.Sleep(15 * time.Second)
+		if unmarshalledData.SleepTime == nil {
+			sleepTime := 15 * time.Second
+			unmarshalledData.SleepTime = &sleepTime
+		}
+
+		if unmarshalledData.MemoryAllocationSize == nil {
+			size := 0
+			unmarshalledData.MemoryAllocationSize = &size
+		}
+
+		fakeAllocation := make([]byte, *unmarshalledData.MemoryAllocationSize)
+		for i := 0; i < *unmarshalledData.MemoryAllocationSize; i++ {
+			fakeAllocation[i] = 0x99
+		}
+
+		time.Sleep(*unmarshalledData.SleepTime)
 		fmt.Println(unmarshalledData)
 
 		ctx.Success()