From 96cb656e5d55a40536edcb8edf14f41026bc9a1c Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Mon, 27 Mar 2023 11:49:35 +0200
Subject: [PATCH] device_write_bps can be int or string with bytes unit

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 loader/loader_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 types/types.go        |  2 +-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/loader/loader_test.go b/loader/loader_test.go
index 2956d7e9..59902af4 100644
--- a/loader/loader_test.go
+++ b/loader/loader_test.go
@@ -2297,3 +2297,43 @@ services:
 	assert.Equal(t, project.Services[0].Name, "extension")
 	assert.Equal(t, project.Services[0].Extensions["x-foo"], "bar")
 }
+
+func TestDeviceWriteBps(t *testing.T) {
+	p, err := loadYAML(`
+        name: test
+        services:
+          foo:
+            image: busybox
+            blkio_config:
+              device_read_bps:
+              - path: /dev/test
+                rate: 1024k
+              device_write_bps:
+              - path: /dev/test
+                rate: 1024
+`)
+	assert.NilError(t, err)
+	assert.DeepEqual(t, p.Services, types.Services{
+		{
+			Name:        "foo",
+			Image:       "busybox",
+			Environment: types.MappingWithEquals{},
+			Scale:       1,
+			BlkioConfig: &types.BlkioConfig{
+				DeviceReadBps: []types.ThrottleDevice{
+					{
+						Path: "/dev/test",
+						Rate: types.UnitBytes(1024 * 1024),
+					},
+				},
+				DeviceWriteBps: []types.ThrottleDevice{
+					{
+						Path: "/dev/test",
+						Rate: types.UnitBytes(1024),
+					},
+				},
+			},
+		},
+	})
+
+}
diff --git a/types/types.go b/types/types.go
index f7399596..c59b540f 100644
--- a/types/types.go
+++ b/types/types.go
@@ -340,7 +340,7 @@ type WeightDevice struct {
 // ThrottleDevice is a structure that holds device:rate_per_second pair
 type ThrottleDevice struct {
 	Path string
-	Rate uint64
+	Rate UnitBytes
 
 	Extensions Extensions `mapstructure:"#extensions" yaml:",inline" json:"-"`
 }