gpu_test.go 858 Bytes
Newer Older
mashun1's avatar
v1  
mashun1 committed
1
2
3
4
5
6
7
package gpu

import (
	"runtime"
	"testing"

	"github.com/stretchr/testify/assert"
xuxzh1's avatar
init  
xuxzh1 committed
8
	"github.com/stretchr/testify/require"
mashun1's avatar
v1  
mashun1 committed
9
10
11
12
)

func TestBasicGetGPUInfo(t *testing.T) {
	info := GetGPUInfo()
xuxzh1's avatar
init  
xuxzh1 committed
13
	assert.NotEmpty(t, len(info))
mashun1's avatar
v1  
mashun1 committed
14
15
16
17
18
19
20
21
22
	assert.Contains(t, "cuda rocm cpu metal", info[0].Library)
	if info[0].Library != "cpu" {
		assert.Greater(t, info[0].TotalMemory, uint64(0))
		assert.Greater(t, info[0].FreeMemory, uint64(0))
	}
}

func TestCPUMemInfo(t *testing.T) {
	info, err := GetCPUMem()
xuxzh1's avatar
init  
xuxzh1 committed
23
	require.NoError(t, err)
mashun1's avatar
v1  
mashun1 committed
24
25
26
27
28
29
30
31
32
33
34
35
	switch runtime.GOOS {
	case "darwin":
		t.Skip("CPU memory not populated on darwin")
	case "linux", "windows":
		assert.Greater(t, info.TotalMemory, uint64(0))
		assert.Greater(t, info.FreeMemory, uint64(0))
	default:
		return
	}
}

// TODO - add some logic to figure out card type through other means and actually verify we got back what we expected