From 56f1744c7ae1b525cff105256afb5dac63a3791a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Osipiuk?= <lukasz@osipiuk.net>
Date: Thu, 26 Apr 2018 17:39:34 +0200
Subject: [PATCH] Add way to request GPUs in tests via RCConfig

---
 test/utils/runners.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/test/utils/runners.go b/test/utils/runners.go
index 28723843862..0a614abb62e 100644
--- a/test/utils/runners.go
+++ b/test/utils/runners.go
@@ -124,6 +124,7 @@ type RCConfig struct {
 	CpuLimit          int64 // millicores
 	MemRequest        int64 // bytes
 	MemLimit          int64 // bytes
+	GpuLimit          int64 // count
 	ReadinessProbe    *v1.Probe
 	DNSPolicy         *v1.DNSPolicy
 	PriorityClassName string
@@ -615,7 +616,7 @@ func (config *RCConfig) applyTo(template *v1.PodTemplateSpec) {
 			c.Ports = append(c.Ports, v1.ContainerPort{Name: k, ContainerPort: int32(v), HostPort: int32(v)})
 		}
 	}
-	if config.CpuLimit > 0 || config.MemLimit > 0 {
+	if config.CpuLimit > 0 || config.MemLimit > 0 || config.GpuLimit > 0 {
 		template.Spec.Containers[0].Resources.Limits = v1.ResourceList{}
 	}
 	if config.CpuLimit > 0 {
@@ -633,6 +634,9 @@ func (config *RCConfig) applyTo(template *v1.PodTemplateSpec) {
 	if config.MemRequest > 0 {
 		template.Spec.Containers[0].Resources.Requests[v1.ResourceMemory] = *resource.NewQuantity(config.MemRequest, resource.DecimalSI)
 	}
+	if config.GpuLimit > 0 {
+		template.Spec.Containers[0].Resources.Limits["nvidia.com/gpu"] = *resource.NewQuantity(config.GpuLimit, resource.DecimalSI)
+	}
 	if len(config.Volumes) > 0 {
 		template.Spec.Volumes = config.Volumes
 	}