11/*
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.
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.
1516 */
1617
1718package e2e
1819
1920import (
2021 "context"
21- "flag"
22+ "os"
23+ "path/filepath"
24+ "runtime"
25+ "strconv"
2226 "testing"
2327
2428 . "github.com/onsi/ginkgo/v2"
@@ -31,33 +35,105 @@ var (
3135
3236 installCTK bool
3337
34- image string
38+ ImageRepo string
39+ ImageTag string
3540
36- sshKey string
37- sshUser string
38- host string
39- sshPort string
40- )
41+ sshKey string
42+ sshUser string
43+ host string
44+ sshPort string
45+ cwd string
46+ packagePath string
4147
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- }
48+ runner Runner
49+ )
5050
5151func TestMain (t * testing.T ) {
52- suiteName := "NVIDIA Container Toolkit E2E "
52+ suiteName := "E2E NVIDIA Container Toolkit"
5353
5454 RegisterFailHandler (Fail )
55+
56+ ctx = context .Background ()
57+ getTestEnv ()
58+
5559 RunSpecs (t ,
5660 suiteName ,
5761 )
5862}
5963
6064// BeforeSuite runs before the test suite
6165var _ = BeforeSuite (func () {
62- ctx = context .Background ()
66+ runner = NewRunner (
67+ WithHost (host ),
68+ WithPort (sshPort ),
69+ WithSshKey (sshKey ),
70+ WithSshUser (sshUser ),
71+ )
72+
73+ if installCTK {
74+ installer , err := NewToolkitInstaller (
75+ WithRunner (runner ),
76+ WithImage (ImageRepo + ":" + ImageTag ),
77+ WithTemplate (dockerInstallTemplate ),
78+ )
79+ Expect (err ).ToNot (HaveOccurred ())
80+
81+ err = installer .Install ()
82+ Expect (err ).ToNot (HaveOccurred ())
83+ }
84+
85+ _ , _ , err := runner .Run ("docker pull ubuntu" )
86+ Expect (err ).ToNot (HaveOccurred ())
87+
88+ _ , _ , err = runner .Run ("docker pull nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda12.5.0" )
89+ Expect (err ).ToNot (HaveOccurred ())
90+
91+ _ , _ , err = runner .Run ("docker pull nvcr.io/nvidia/cuda:12.8.0-base-ubi8" )
92+ Expect (err ).ToNot (HaveOccurred ())
6393})
94+
95+ // getTestEnv gets the test environment variables
96+ func getTestEnv () {
97+ defer GinkgoRecover ()
98+ var err error
99+
100+ _ , thisFile , _ , _ := runtime .Caller (0 )
101+ packagePath = filepath .Dir (thisFile )
102+
103+ installCTK = getBoolEnvVar ("INSTALL_CTK" , false )
104+
105+ ImageRepo = os .Getenv ("E2E_IMAGE_REPO" )
106+ Expect (ImageRepo ).NotTo (BeEmpty (), "E2E_IMAGE_REPO environment variable must be set" )
107+
108+ ImageTag = os .Getenv ("E2E_IMAGE_TAG" )
109+ Expect (ImageTag ).NotTo (BeEmpty (), "E2E_IMAGE_TAG environment variable must be set" )
110+
111+ sshKey = os .Getenv ("SSH_KEY" )
112+ Expect (sshKey ).NotTo (BeEmpty (), "SSH_KEY environment variable must be set" )
113+
114+ sshUser = os .Getenv ("SSH_USER" )
115+ Expect (sshUser ).NotTo (BeEmpty (), "SSH_USER environment variable must be set" )
116+
117+ host = os .Getenv ("REMOTE_HOST" )
118+ Expect (host ).NotTo (BeEmpty (), "REMOTE_HOST environment variable must be set" )
119+
120+ sshPort = os .Getenv ("REMOTE_PORT" )
121+ Expect (sshPort ).NotTo (BeEmpty (), "REMOTE_PORT environment variable must be set" )
122+
123+ // Get current working directory
124+ cwd , err = os .Getwd ()
125+ Expect (err ).NotTo (HaveOccurred ())
126+ }
127+
128+ // getBoolEnvVar returns the boolean value of the environment variable or the default value if not set.
129+ func getBoolEnvVar (key string , defaultValue bool ) bool {
130+ value := os .Getenv (key )
131+ if value == "" {
132+ return defaultValue
133+ }
134+ boolValue , err := strconv .ParseBool (value )
135+ if err != nil {
136+ return defaultValue
137+ }
138+ return boolValue
139+ }
0 commit comments