mirror of https://github.com/portainer/portainer
76 lines
1.9 KiB
Go
76 lines
1.9 KiB
Go
package postinit
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
portainer "github.com/portainer/portainer/api"
|
|
"github.com/portainer/portainer/api/datastore"
|
|
|
|
"github.com/docker/docker/api/types/container"
|
|
"github.com/docker/docker/client"
|
|
"github.com/segmentio/encoding/json"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestMigrateGPUs(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if strings.HasSuffix(r.URL.Path, "/containers/json") {
|
|
containerSummary := []container.Summary{{ID: "container1"}}
|
|
|
|
err := json.NewEncoder(w).Encode(containerSummary)
|
|
require.NoError(t, err)
|
|
|
|
return
|
|
}
|
|
|
|
container := container.InspectResponse{
|
|
ContainerJSONBase: &container.ContainerJSONBase{
|
|
ID: "container1",
|
|
HostConfig: &container.HostConfig{
|
|
Resources: container.Resources{
|
|
DeviceRequests: []container.DeviceRequest{
|
|
{Driver: "nvidia"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
err := json.NewEncoder(w).Encode(container)
|
|
require.NoError(t, err)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
_, store := datastore.MustNewTestStore(t, true, false)
|
|
|
|
migrator := &PostInitMigrator{dataStore: store}
|
|
|
|
dockerCli, err := client.NewClientWithOpts(client.WithHost(srv.URL), client.WithHTTPClient(http.DefaultClient))
|
|
require.NoError(t, err)
|
|
|
|
// Nonexistent endpoint
|
|
|
|
err = migrator.MigrateGPUs(portainer.Endpoint{}, dockerCli)
|
|
require.Error(t, err)
|
|
|
|
// Valid endpoint
|
|
|
|
endpoint := portainer.Endpoint{ID: 1, PostInitMigrations: portainer.EndpointPostInitMigrations{MigrateGPUs: true}}
|
|
|
|
err = store.Endpoint().Create(&endpoint)
|
|
require.NoError(t, err)
|
|
|
|
err = migrator.MigrateGPUs(endpoint, dockerCli)
|
|
require.NoError(t, err)
|
|
|
|
migratedEndpoint, err := store.Endpoint().Endpoint(endpoint.ID)
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, endpoint.ID, migratedEndpoint.ID)
|
|
require.False(t, migratedEndpoint.PostInitMigrations.MigrateGPUs)
|
|
require.True(t, migratedEndpoint.EnableGPUManagement)
|
|
}
|