gpu_test.go 858 Bytes
Newer Older
1
2
3
4
5
6
7
package gpu

import (
	"runtime"
	"testing"

	"github.com/stretchr/testify/assert"
Michael Yang's avatar
lint  
Michael Yang committed
8
	"github.com/stretchr/testify/require"
9
10
11
12
)

func TestBasicGetGPUInfo(t *testing.T) {
	info := GetGPUInfo()
Michael Yang's avatar
lint  
Michael Yang committed
13
	assert.NotEmpty(t, len(info))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
14
15
16
17
	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))
18
19
20
	}
}

21
func TestCPUMemInfo(t *testing.T) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
22
	info, err := GetCPUMem()
Michael Yang's avatar
lint  
Michael Yang committed
23
	require.NoError(t, err)
24
25
26
27
28
29
30
31
32
33
34
	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
	}
}

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