| package gpu | |
| import ( | |
| "bytes" | |
| "fmt" | |
| "io" | |
| "net/http" | |
| "os" | |
| "path/filepath" | |
| "strings" | |
| "testing" | |
| "github.com/google/go-cmp/cmp" | |
| ) | |
| type FakeHttpClient struct { | |
| URL string | |
| Header http.Header | |
| Body []byte | |
| } | |
| func (h *FakeHttpClient) Do(r *http.Request) (*http.Response, error) { | |
| h.URL = r.URL.String() | |
| h.Header = r.Header | |
| resp := &http.Response{ | |
| Header: r.Header, | |
| StatusCode: http.StatusOK, | |
| Request: r, | |
| Body: io.NopCloser(bytes.NewReader(h.Body)), | |
| } | |
| return resp, nil | |
| } | |
| type FakeOSUtils struct { | |
| result string | |
| cmds []string | |
| } | |
| func (f *FakeOSUtils) run(command string, args []string, hideStderr bool) (string, error) { | |
| if command == "uname" { | |
| return f.result, nil | |
| } | |
| cmd := fmt.Sprintf("%s %s", command, strings.Join(args, " ")) | |
| f.cmds = append(f.cmds, command) | |
| f.cmds = append(f.cmds, args...) | |
| // Returns the last command ran. | |
| return cmd, nil | |
| } | |
| func TestRetrieveInstanceZone(t *testing.T) { | |
| h := &FakeHttpClient{Body: []byte(`projects/475556798229/zones/us-central1-a`)} | |
| zone, err := retrieveInstanceZone(h, instanceURL) | |
| if err != nil { | |
| t.Fatalf("retrieveInstanceZone failed: %v", err) | |
| } | |
| wantZone := "us" | |
| wantHeader := "Google" | |
| if h.URL != instanceURL || h.Header["Metadata-Flavor"][0] != wantHeader { | |
| t.Errorf("TestRetrieveInstanceZone failed: Request not valid:\n wantUrl: %s\tgotUrl:%s\n wantHeader:%v\n gotHeader:%v", | |
| h.URL, instanceURL, "Google", h.Header["Metadata-Flavor"][0]) | |
| } | |
| if zone != wantZone { | |
| t.Errorf("TestRetrieveInstanceZone failed: Returned unexpected difference:\n wantUrl: %s\tgotUrl:%s\n wantHeader:%v\n gotHeader:%v", | |
| h.URL, instanceURL, "Google", h.Header["Metadata-Flavor"][0]) | |
| } | |
| } | |
| func TestCheckArch(t *testing.T) { | |
| tests := []struct { | |
| desc string | |
| wantErr bool | |
| hideStderr bool | |
| arch string | |
| }{ | |
| { | |
| desc: "Success when arch is x86_64", | |
| wantErr: false, | |
| hideStderr: false, | |
| arch: "x86_64", | |
| }, | |
| { | |
| desc: "Failure when arch is ARM", | |
| wantErr: false, | |
| hideStderr: false, | |
| arch: "aarch64", | |
| }, | |
| { | |
| desc: "Failure when arch is other", | |
| wantErr: true, | |
| hideStderr: false, | |
| arch: "ARM", | |
| }, | |
| } | |
| for _, test := range tests { | |
| t.Run(test.desc, func(t *testing.T) { | |
| f := FakeOSUtils{result: test.arch} | |
| run = f.run | |
| err := checkArch() | |
| if gotErr := err != nil; gotErr != test.wantErr { | |
| t.Errorf("TestCheckArch(%s): Error: %s\n gotErr: %t\n wantErr: %t", test.desc, err, gotErr, test.wantErr) | |
| } | |
| }) | |
| } | |
| } | |
| func TestGetCosInstaller(t *testing.T) { | |
| tests := []struct { | |
| desc string | |
| wantErr bool | |
| httpRes []byte | |
| wantInstaller Installers | |
| setEnv bool | |
| setInstaller string | |
| apiDomains string | |
| }{ | |
| { | |
| desc: "Vm instance zone in artifact regions", | |
| wantErr: false, | |
| httpRes: []byte("us"), | |
| wantInstaller: Installers{DefaultInstaller: "us.gcr.io/cos-cloud/cos-gpu-installer:v2.4.4", SelectedInstaller: ""}, | |
| }, | |
| { | |
| desc: "Vm instance zone not in artifact regions", | |
| wantErr: false, | |
| httpRes: []byte("cn"), | |
| wantInstaller: Installers{DefaultInstaller: "gcr.io/cos-cloud/cos-gpu-installer:v2.4.4", SelectedInstaller: ""}, | |
| }, | |
| { | |
| desc: "Installer set in env variables.", | |
| wantErr: false, | |
| setEnv: true, | |
| setInstaller: "cloud-installer-101", | |
| httpRes: []byte("us"), | |
| wantInstaller: Installers{DefaultInstaller: "gcr.io/cos-cloud/cos-gpu-installer:v2.4.4", SelectedInstaller: "cloud-installer-101"}, | |
| }, | |
| { | |
| desc: "Running in TPC", | |
| wantErr: false, | |
| httpRes: []byte("u"), | |
| wantInstaller: Installers{DefaultInstaller: "docker.fake-ar-domain.goog/fake-project-prefix/cos-cloud/cos-gpu-installer/cos-gpu-installer:v2.4.4", SelectedInstaller: ""}, | |
| apiDomains: "API_DOMAIN=fake-api-domain.goog\nARTIFACT_REGISTRY_DOMAIN=fake-ar-domain.goog\nPROJECT_PREFIX=fake-project-prefix", | |
| }, | |
| { | |
| desc: "Running in GDU with domain file", | |
| wantErr: false, | |
| httpRes: []byte("us"), | |
| wantInstaller: Installers{DefaultInstaller: "us.gcr.io/cos-cloud/cos-gpu-installer:v2.4.4", SelectedInstaller: ""}, | |
| apiDomains: "API_DOMAIN=googleapis.com\nARTIFACT_REGISTRY_DOMAIN=pkg.dev\nPROJECT_PREFIX=", | |
| }, | |
| } | |
| for _, test := range tests { | |
| t.Run(test.desc, func(t *testing.T) { | |
| c := &FakeHttpClient{Body: test.httpRes} | |
| if test.apiDomains != "" { | |
| tmpFile, err := os.CreateTemp("", "") | |
| if err != nil { | |
| t.Fatalf("failed to create tmpFile: %v", err) | |
| } | |
| defer os.Remove(tmpFile.Name()) | |
| _, err = tmpFile.WriteString(test.apiDomains) | |
| if err != nil { | |
| t.Fatalf("failed to write to tmpFile: %v", err) | |
| } | |
| realAPIDomainsPath := apiDomainsPath | |
| apiDomainsPath = tmpFile.Name() | |
| defer func() { apiDomainsPath = realAPIDomainsPath }() | |
| } | |
| if test.setEnv == true { | |
| t.Setenv("COS_GPU_INSTALLER", test.setInstaller) | |
| } | |
| cosInstaller, err := GetCosInstaller(c) | |
| if gotErr := err != nil; gotErr != test.wantErr { | |
| t.Errorf("TestGetCosInstaller(%s): Error: %s\n gotErr: %t\n wantErr: %t", test.desc, err, gotErr, test.wantErr) | |
| } | |
| if diff := cmp.Diff(cosInstaller, test.wantInstaller); diff != "" { | |
| t.Errorf("TestGetCosInstaller(%s): returned unexpected difference;\n wantInstaller:%s\t gotInstaller:%s", test.desc, test.wantInstaller, cosInstaller) | |
| } | |
| }) | |
| } | |
| } | |
| func TestList(t *testing.T) { | |
| tests := []struct { | |
| desc string | |
| args []string | |
| wantErr bool | |
| listInstaller bool | |
| installer Installers | |
| want string | |
| }{ | |
| { | |
| desc: "Flag --gpu-installer passed in", | |
| args: []string{}, | |
| listInstaller: true, | |
| installer: Installers{DefaultInstaller: "installer", SelectedInstaller: "selected installer"}, | |
| want: "installer", | |
| }, | |
| { | |
| desc: "Flag --gpu-installer not passed in", | |
| args: []string{}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| want: "Available", | |
| }, | |
| } | |
| for _, test := range tests { | |
| t.Run(test.desc, func(t *testing.T) { | |
| var got bytes.Buffer | |
| writer = &got | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| err := List(test.installer, test.listInstaller, test.args) | |
| if gotErr := err != nil; gotErr != test.wantErr { | |
| t.Errorf("TestList(%s): Error: %s\n gotErr: %t\n wantErr: %t", test.desc, err, gotErr, test.wantErr) | |
| } | |
| if !strings.HasPrefix(got.String(), test.want) { | |
| t.Errorf("TestGetCosInstaller(%s): returned unexpected difference;\n want:%s\t got:%s", test.desc, test.want, got.String()) | |
| } | |
| }) | |
| } | |
| } | |
| func TestRunInstaller(t *testing.T) { | |
| tests := []struct { | |
| desc string | |
| wantArgs []string | |
| cacheExists bool | |
| wantErr bool | |
| installer Installers | |
| cmd string | |
| args []string | |
| hideStderr bool | |
| }{ | |
| { | |
| desc: "Uses set installer", | |
| wantArgs: []string{ | |
| "/usr/bin/docker", | |
| "run", | |
| "--rm", | |
| "--name=cos-gpu-installer", | |
| "--privileged", | |
| "--net=host", | |
| "--pid=host", | |
| "--volume", "/dev:/dev", | |
| "--volume", "/:/root", | |
| "--log-driver", "journald", | |
| "installer", | |
| "install", | |
| }, | |
| wantErr: false, | |
| cmd: "install", | |
| installer: Installers{SelectedInstaller: "installer", DefaultInstaller: "default installer"}, | |
| hideStderr: false, | |
| }, | |
| { | |
| desc: "Uses extension cache: flag --prepare-build-tools passed in", | |
| wantErr: false, | |
| args: []string{"--prepare-build-tools"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| cmd: "install", | |
| wantArgs: []string{ | |
| "/usr/bin/docker", | |
| "run", | |
| "--rm", | |
| "--name=cos-gpu-installer", | |
| "--privileged", | |
| "--net=host", | |
| "--pid=host", | |
| "--volume", "/dev:/dev", | |
| "--volume", "/:/root", | |
| "--log-driver", "journald", | |
| "--volume", "/var/lib/cos-extensions/:/build/", | |
| "installer", | |
| "install", | |
| "--prepare-build-tools", | |
| }, | |
| hideStderr: false, | |
| }, | |
| { | |
| desc: "Does not use extension cache", | |
| wantArgs: []string{ | |
| "/usr/bin/docker", | |
| "run", | |
| "--rm", | |
| "--name=cos-gpu-installer", | |
| "--privileged", | |
| "--net=host", | |
| "--pid=host", | |
| "--volume", "/dev:/dev", | |
| "--volume", "/:/root", | |
| "--log-driver", "journald", | |
| "installer", | |
| "install", | |
| }, | |
| wantErr: false, | |
| cmd: "install", | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| hideStderr: false, | |
| }, | |
| } | |
| for _, test := range tests { | |
| t.Run(test.desc, func(t *testing.T) { | |
| var got bytes.Buffer | |
| writer = &got | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| err := runInstaller(test.cmd, test.hideStderr, extensionCache, test.installer, test.args...) | |
| if gotErr := err != nil; gotErr != test.wantErr { | |
| t.Errorf("TestList(%s): Error: %s\n gotErr: %t\n wantErr: %t", test.desc, err, gotErr, test.wantErr) | |
| } | |
| if !cmp.Equal(f.cmds, test.wantArgs) { | |
| t.Errorf("TestRunInstaller(%s): returned unexpected difference;\n want:%s\t got:%v", test.desc, test.wantArgs, got) | |
| } | |
| }) | |
| } | |
| } | |
| func TestRunInstallerCacheExists(t *testing.T) { | |
| tmpFile, err := os.CreateTemp(".", "cos-extensions") | |
| if err != nil { | |
| t.Fatalf("Failed to create temp file: %v", err) | |
| } | |
| cacheFile := tmpFile.Name() | |
| defer os.Remove(tmpFile.Name()) | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| err = runInstaller("install", true, cacheFile, Installers{DefaultInstaller: "installer"}) | |
| if gotErr := err != nil; gotErr { | |
| t.Errorf("TestRunInstallerCacheExists: Error: %s\n", err) | |
| } | |
| wantArgs := []string{ | |
| "/usr/bin/docker", | |
| "run", | |
| "--rm", | |
| "--name=cos-gpu-installer", | |
| "--privileged", | |
| "--net=host", | |
| "--pid=host", | |
| "--volume", "/dev:/dev", | |
| "--volume", "/:/root", | |
| "--log-driver", "journald", | |
| "--volume", fmt.Sprintf("%s/:/build/", tmpFile.Name()), | |
| "installer", | |
| "install", | |
| } | |
| if !cmp.Equal(f.cmds, wantArgs) { | |
| t.Errorf("TestRunInstallerCacheExists: returned unexpected difference;\n want:%v\t got:%v", wantArgs, f.cmds) | |
| } | |
| } | |
| func TestRunInstallerCleanCache(t *testing.T) { | |
| tmpDir, err := os.MkdirTemp("", "tmpDir") | |
| if err != nil { | |
| t.Fatalf("Failed to create temp file: %v", err) | |
| } | |
| args := []string{"--clean-build-tools"} | |
| // Ensuring temp dir is deleted if test fails. | |
| defer os.RemoveAll(tmpDir) | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| err = runInstaller("install", false, tmpDir, Installers{DefaultInstaller: "installer"}, args...) | |
| if err != nil { | |
| t.Errorf("TestRunInstallerCleanCache: Error: %s\n", err) | |
| } | |
| if _, err := os.Stat(tmpDir); !os.IsNotExist(err) { | |
| t.Errorf("TestRunInstallerCleanCache: did not remove cache file") | |
| } | |
| } | |
| func TestInstall(t *testing.T) { | |
| tests := []struct { | |
| desc string | |
| args []string | |
| wantCmds []string | |
| installer Installers | |
| wantOutput string | |
| }{ | |
| { | |
| desc: "Default hostDir", | |
| args: []string{}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "-host-dir=/var/lib/nvidia", | |
| }, | |
| }, | |
| { | |
| desc: "Custom hostDir", | |
| args: []string{"-host-dir=/var/lib/tmp"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "-host-dir=/var/lib/tmp", | |
| }, | |
| }, | |
| { | |
| desc: "No verify", | |
| args: []string{"-no-verify"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "-host-dir=/var/lib/nvidia", "-no-verify", | |
| }, | |
| wantOutput: "Skipping post_install due to no-verify flag.\n", | |
| }, | |
| { | |
| desc: "Space separated --host-dir", | |
| args: []string{"--host-dir", "/var/lib/tmp"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "--host-dir", "/var/lib/tmp", | |
| }, | |
| }, | |
| { | |
| desc: "Space separated -host-dir", | |
| args: []string{"-host-dir", "/var/lib/tmp"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "-host-dir", "/var/lib/tmp", | |
| }, | |
| }, | |
| { | |
| desc: "Flags at the end without argument", | |
| args: []string{"-host-dir"}, | |
| installer: Installers{DefaultInstaller: "installer"}, | |
| wantCmds: []string{ | |
| "/usr/bin/docker", "run", "--rm", "--name=cos-gpu-installer", | |
| "--privileged", "--net=host", "--pid=host", "--volume", "/dev:/dev", | |
| "--volume", "/:/root", "--log-driver", "journald", "installer", | |
| "install", "-host-dir=/var/lib/nvidia", "-host-dir", | |
| }, | |
| }, | |
| } | |
| for _, test := range tests { | |
| t.Run(test.desc, func(t *testing.T) { | |
| var got bytes.Buffer | |
| writer = &got | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| err := Install(test.installer, test.args) | |
| if err != nil { | |
| t.Errorf("TestInstall(%s): Error: %s", test.desc, err) | |
| } | |
| want := strings.Join(test.wantCmds, " ") | |
| if got := strings.Join(f.cmds, " "); !strings.HasPrefix(got, want) { | |
| t.Errorf("TestInstall(%s): returned unexpected difference in commands;\n want:%s\t got:%s", test.desc, want, got) | |
| } | |
| if test.wantOutput != "" && !strings.HasSuffix(got.String(), test.wantOutput) { | |
| t.Errorf("TestInstall(%s): returned unexpected difference in output;\n want:%s\t got:%s", test.desc, test.wantOutput, got.String()) | |
| } | |
| }) | |
| } | |
| } | |
| func TestPostInstall(t *testing.T) { | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| hostDir := t.TempDir() | |
| postInstall(hostDir) | |
| wantCmds := []string{ | |
| "/bin/sh", "-c", | |
| fmt.Sprintf("sudo -n -- sh -c \"mount --bind %[1]s %[1]s; "+ | |
| "mount -o remount,exec %[1]s\"", hostDir), | |
| "/bin/sh", "-c", | |
| fmt.Sprintf("sudo %s/bin/nvidia-persistenced", hostDir), | |
| } | |
| if !cmp.Equal(f.cmds, wantCmds) { | |
| t.Errorf("TestPostInstall: returned unexpected difference in commands;\n want:%v\t got:%v", wantCmds, f.cmds) | |
| } | |
| } | |
| func TestPostInstallWithGridd(t *testing.T) { | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| hostDir := t.TempDir() | |
| binDir := filepath.Join(hostDir, "bin") | |
| if err := os.MkdirAll(binDir, 0755); err != nil { | |
| t.Fatalf("failed to create bin dir: %v", err) | |
| } | |
| if err := os.WriteFile(filepath.Join(binDir, "nvidia-gridd"), []byte("fake"), 0755); err != nil { | |
| t.Fatalf("failed to create fake nvidia-gridd: %v", err) | |
| } | |
| fakeService := filepath.Join(hostDir, "nvidia-persistenced.service") | |
| if err := os.WriteFile(fakeService, []byte("[Unit]\n"), 0644); err != nil { | |
| t.Fatalf("failed to create fake nvidia-persistenced.service: %v", err) | |
| } | |
| origServicePath := persistencedServicePath | |
| persistencedServicePath = fakeService | |
| defer func() { persistencedServicePath = origServicePath }() | |
| postInstall(hostDir) | |
| if len(f.cmds) != 0 { | |
| t.Errorf("TestPostInstallWithGridd: expected no commands to be run when nvidia-gridd and nvidia-persistenced.service are present, got: %v", f.cmds) | |
| } | |
| } | |
| func TestPostInstallWithGriddNoPersistencedService(t *testing.T) { | |
| f := FakeOSUtils{result: "x86_64"} | |
| run = f.run | |
| hostDir := t.TempDir() | |
| binDir := filepath.Join(hostDir, "bin") | |
| if err := os.MkdirAll(binDir, 0755); err != nil { | |
| t.Fatalf("failed to create bin dir: %v", err) | |
| } | |
| if err := os.WriteFile(filepath.Join(binDir, "nvidia-gridd"), []byte("fake"), 0755); err != nil { | |
| t.Fatalf("failed to create fake nvidia-gridd: %v", err) | |
| } | |
| origServicePath := persistencedServicePath | |
| persistencedServicePath = filepath.Join(hostDir, "nonexistent-nvidia-persistenced.service") | |
| defer func() { persistencedServicePath = origServicePath }() | |
| postInstall(hostDir) | |
| wantCmds := []string{ | |
| "/bin/sh", "-c", | |
| fmt.Sprintf("sudo -n -- sh -c \"mount --bind %[1]s %[1]s; "+ | |
| "mount -o remount,exec %[1]s\"", hostDir), | |
| "/bin/sh", "-c", | |
| fmt.Sprintf("sudo %s/bin/nvidia-persistenced", hostDir), | |
| } | |
| if !cmp.Equal(f.cmds, wantCmds) { | |
| t.Errorf("TestPostInstallWithGriddNoPersistencedService: returned unexpected difference in commands;\n want:%v\t got:%v", wantCmds, f.cmds) | |
| } | |
| } |