NVIDIA GPU

/*

** Zabbix

** Copyright (C) 2001-2024 Zabbix SIA

**

** Licensed under the Apache License, Version 2.0 (the "License");

** you may not use this file except in compliance with the License.

** You may obtain a copy of the License at

**

**     http://www.apache.org/licenses/LICENSE-2.0

**

** Unless required by applicable law or agreed to in writing, software

** distributed under the License is distributed on an "AS IS" BASIS,

** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

** See the License for the specific language governing permissions and

** limitations under the License.

**/

​

package nvmlmock

​

import (

    "testing"

​

    "github.com/google/go-cmp/cmp"

    "golang.zabbix.com/plugin/nvidia/pkg/nvml"

)

​

var (

    _ nvml.Device = (*MockDevice)(nil)

    _ Mocker      = (*MockDevice)(nil)

)

​

// MockDevice is mock for NVML device.

type MockDevice struct {

    nvml.Device

    expectations []*Expectation

    callIdx      int

    t            *testing.T

}

​

// NewMockDevice returns new mock device.

func NewMockDevice(t *testing.T) *MockDevice {

    t.Helper()

​

    return &MockDevice{

        t:            t,

        expectations: []*Expectation{},

    }

}

​

// ExpectCalls sets calls that are expected by mock.

func (m *MockDevice) ExpectCalls(expectations ...*Expectation) *MockDevice {

    m.expectations = expectations

​

    return m

}

​

// ExpectedCallsDone checks if all expected calls of mock and it's submocks are done.

func (m *MockDevice) ExpectedCallsDone() bool {

    m.t.Helper()

    expected := len(m.expectations)

    received := m.callIdx

​

    if expected == received {

        return true

    }

​

    for _, e := range m.expectations[received:] {

        m.t.Errorf("Not called %q", e.funcName)

    }

​

    m.t.Errorf("received %d out of %d expected calls", received, expected)

​

    return false

}

​

// SubMocks returns submocks of the mock.

func (m *MockDevice) SubMocks() []Mocker {

    var subMocks []Mocker

​

    for _, expect := range m.expectations {

        for _, out := range expect.out {

            subMock, ok := out.(Mocker)

            if !ok {

                continue

            }

​

            subMocks = append(subMocks, subMock)

            subMocks = append(subMocks, subMock.SubMocks()...)

        }

    }

​