gpu_test.go 581 B

1234567891011121314151617181920212223242526
  1. package gpu
  2. import (
  3. "runtime"
  4. "testing"
  5. "github.com/stretchr/testify/assert"
  6. )
  7. func TestBasicGetGPUInfo(t *testing.T) {
  8. info := GetGPUInfo()
  9. assert.Contains(t, "CUDA ROCM CPU METAL", info.Driver)
  10. switch runtime.GOOS {
  11. case "darwin":
  12. // TODO - remove this once MacOS returns some size for CPU
  13. return
  14. case "linux", "windows":
  15. assert.Greater(t, info.TotalMemory, uint64(0))
  16. assert.Greater(t, info.FreeMemory, uint64(0))
  17. default:
  18. return
  19. }
  20. }
  21. // TODO - add some logic to figure out card type through other means and actually verify we got back what we expected