iptablesprovidermock.go

package provider

import (
	"sync"
	"testing"
)

type iptablesProviderMockedMethods struct {
	appendMock        func(table, chain string, rulespec ...string) error
	insertMock        func(table, chain string, pos int, rulespec ...string) error
	deleteMock        func(table, chain string, rulespec ...string) error
	listChainsMock    func(table string) ([]string, error)
	clearChainMock    func(table, chain string) error
	deleteChainMock   func(table, chain string) error
	newChainMock      func(table, chain string) error
	commitMock        func() error
	retrieveTableMock func() map[string]map[string][]string
}

// TestIptablesProvider is a test implementation for IptablesProvider
type TestIptablesProvider interface {
	IptablesProvider
	MockAppend(t *testing.T, impl func(table, chain string, rulespec ...string) error)
	MockInsert(t *testing.T, impl func(table, chain string, pos int, rulespec ...string) error)
	MockDelete(t *testing.T, impl func(table, chain string, rulespec ...string) error)
	MockListChains(t *testing.T, impl func(table string) ([]string, error))
	MockClearChain(t *testing.T, impl func(table, chain string) error)
	MockDeleteChain(t *testing.T, impl func(table, chain string) error)
	MockNewChain(t *testing.T, impl func(table, chain string) error)
	MockCommit(t *testing.T, impl func() error)
}

// A testIptablesProvider is an empty TransactionalManipulator that can be easily mocked.
type testIptablesProvider struct {
	mocks       map[*testing.T]*iptablesProviderMockedMethods
	lock        *sync.Mutex
	currentTest *testing.T
}

// NewTestIptablesProvider returns a new TestManipulator.
func NewTestIptablesProvider() TestIptablesProvider {
	return &testIptablesProvider{lock: &sync.Mutex{}, mocks: map[*testing.T]*iptablesProviderMockedMethods{}}
}

func (m *testIptablesProvider) MockAppend(t *testing.T, impl func(table, chain string, rulespec ...string) error) {

	m.currentMocks(t).appendMock = impl
}

func (m *testIptablesProvider) MockInsert(t *testing.T, impl func(table, chain string, pos int, rulespec ...string) error) {

	m.currentMocks(t).insertMock = impl
}

func (m *testIptablesProvider) MockDelete(t *testing.T, impl func(table, chain string, rulespec ...string) error) {

	m.currentMocks(t).deleteMock = impl
}

func (m *testIptablesProvider) MockListChains(t *testing.T, impl func(table string) ([]string, error)) {

	m.currentMocks(t).listChainsMock = impl
}

func (m *testIptablesProvider) MockClearChain(t *testing.T, impl func(table, chain string) error) {

	m.currentMocks(t).clearChainMock = impl
}

func (m *testIptablesProvider) MockDeleteChain(t *testing.T, impl func(table, chain string) error) {

	m.currentMocks(t).deleteChainMock = impl
}

func (m *testIptablesProvider) MockNewChain(t *testing.T, impl func(table, chain string) error) {

	m.currentMocks(t).newChainMock = impl
}

func (m *testIptablesProvider) MockCommit(t *testing.T, impl func() error) {
	m.currentMocks(t).commitMock = impl
}

func (m *testIptablesProvider) Append(table, chain string, rulespec ...string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.appendMock != nil {
		return mock.appendMock(table, chain, rulespec...)
	}

	return nil
}

func (m *testIptablesProvider) Insert(table, chain string, pos int, rulespec ...string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.insertMock != nil {
		return mock.insertMock(table, chain, pos, rulespec...)
	}

	return nil
}

func (m *testIptablesProvider) Delete(table, chain string, rulespec ...string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.deleteMock != nil {
		return mock.deleteMock(table, chain, rulespec...)
	}

	return nil
}

func (m *testIptablesProvider) ListChains(table string) ([]string, error) {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.listChainsMock != nil {
		return mock.listChainsMock(table)
	}

	return nil, nil
}

func (m *testIptablesProvider) ClearChain(table, chain string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.clearChainMock != nil {
		return mock.clearChainMock(table, chain)
	}

	return nil
}

func (m *testIptablesProvider) DeleteChain(table, chain string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.deleteChainMock != nil {
		return mock.deleteChainMock(table, chain)
	}

	return nil
}

func (m *testIptablesProvider) NewChain(table, chain string) error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.newChainMock != nil {
		return mock.newChainMock(table, chain)
	}

	return nil
}

func (m *testIptablesProvider) Commit() error {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.commitMock != nil {
		return mock.commitMock()
	}

	return nil
}

func (m *testIptablesProvider) RetrieveTable() map[string]map[string][]string {

	if mock := m.currentMocks(m.currentTest); mock != nil && mock.retrieveTableMock != nil {
		return mock.retrieveTableMock()
	}

	return nil
}

func (m *testIptablesProvider) currentMocks(t *testing.T) *iptablesProviderMockedMethods {
	m.lock.Lock()
	defer m.lock.Unlock()

	mocks := m.mocks[t]

	if mocks == nil {
		mocks = &iptablesProviderMockedMethods{}
		m.mocks[t] = mocks
	}

	m.currentTest = t
	return mocks
}