1- /*
2- * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3- *
4- * Licensed under the Apache License, Version 2.0 (the "License");
5- * you may not use this file except in compliance with the License.
6- * You may obtain a copy of the License at
7- *
8- * http://www.apache.org/licenses/LICENSE-2.0
9- *
10- * Unless required by applicable law or agreed to in writing, software
11- * distributed under the License is distributed on an "AS IS" BASIS,
12- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13- * See the License for the specific language governing permissions and
14- * limitations under the License.
15- */
1+ /**
2+ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+ # SPDX-License-Identifier: Apache-2.0
4+ #
5+ # Licensed under the Apache License, Version 2.0 (the "License");
6+ # you may not use this file except in compliance with the License.
7+ # You may obtain a copy of the License at
8+ #
9+ # http://www.apache.org/licenses/LICENSE-2.0
10+ #
11+ # Unless required by applicable law or agreed to in writing, software
12+ # distributed under the License is distributed on an "AS IS" BASIS,
13+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+ # See the License for the specific language governing permissions and
15+ # limitations under the License.
16+ **/
1617
1718package e2e
1819
1920import (
2021 "context"
21- "flag"
22+ "errors"
23+ "os"
24+ "strconv"
2225 "testing"
2326
2427 . "github.com/onsi/ginkgo/v2"
@@ -31,33 +34,86 @@ var (
3134
3235 installCTK bool
3336
34- image string
37+ imageName string
38+ imageTag string
3539
3640 sshKey string
3741 sshUser string
38- host string
42+ sshHost string
3943 sshPort string
4044)
4145
42- func init () {
43- flag .BoolVar (& installCTK , "install-ctk" , false , "Install the NVIDIA Container Toolkit" )
44- flag .StringVar (& image , "toolkit-image" , "" , "Repository of the image to test" )
45- flag .StringVar (& sshKey , "ssh-key" , "" , "SSH key to use for remote login" )
46- flag .StringVar (& sshUser , "ssh-user" , "" , "SSH user to use for remote login" )
47- flag .StringVar (& host , "remote-host" , "" , "Hostname of the remote machine" )
48- flag .StringVar (& sshPort , "remote-port" , "22" , "SSH port to use for remote login" )
49- }
50-
5146func TestMain (t * testing.T ) {
52- suiteName := "NVIDIA Container Toolkit E2E "
47+ suiteName := "E2E NVIDIA Container Toolkit"
5348
5449 RegisterFailHandler (Fail )
50+
51+ ctx = context .Background ()
52+ getTestEnv ()
53+
5554 RunSpecs (t ,
5655 suiteName ,
5756 )
5857}
5958
60- // BeforeSuite runs before the test suite
61- var _ = BeforeSuite (func () {
62- ctx = context .Background ()
63- })
59+ // getTestEnv gets the test environment variables
60+ func getTestEnv () {
61+ defer GinkgoRecover ()
62+
63+ installCTK = getEnvVarOrDefault ("E2E_INSTALL_CTK" , false )
64+
65+ if installCTK {
66+ imageName = getRequiredEnvvar [string ]("E2E_IMAGE_NAME" )
67+
68+ imageTag = getRequiredEnvvar [string ]("E2E_IMAGE_TAG" )
69+
70+ }
71+
72+ sshKey = getRequiredEnvvar [string ]("E2E_SSH_KEY" )
73+ sshUser = getRequiredEnvvar [string ]("E2E_SSH_USER" )
74+ sshHost = getRequiredEnvvar [string ]("E2E_SSH_HOST" )
75+
76+ sshPort = getEnvVarOrDefault ("E2E_SSH_PORT" , "22" )
77+ }
78+
79+ // getRequiredEnvvar returns the specified envvar if set or raises an error.
80+ func getRequiredEnvvar [T any ](key string ) T {
81+ v , err := getEnvVarAs [T ](key )
82+ Expect (err ).To (BeNil (), "required environement variable not set" , key )
83+ return v
84+ }
85+
86+ func getEnvVarAs [T any ](key string ) (T , error ) {
87+ var zero T
88+ value := os .Getenv (key )
89+ if value == "" {
90+ return zero , errors .New ("env var not set" )
91+ }
92+
93+ switch any (zero ).(type ) {
94+ case bool :
95+ v , err := strconv .ParseBool (value )
96+ if err != nil {
97+ return zero , err
98+ }
99+ return any (v ).(T ), nil
100+ case int :
101+ v , err := strconv .Atoi (value )
102+ if err != nil {
103+ return zero , err
104+ }
105+ return any (v ).(T ), nil
106+ case string :
107+ return any (value ).(T ), nil
108+ default :
109+ return zero , errors .New ("unsupported type" )
110+ }
111+ }
112+
113+ func getEnvVarOrDefault [T any ](key string , defaultValue T ) T {
114+ val , err := getEnvVarAs [T ](key )
115+ if err != nil {
116+ return defaultValue
117+ }
118+ return val
119+ }
0 commit comments