firewall_containers/network-go/mock/mock.go

package mock

import (
	"context"
	"time"

	"github.com/docker/docker/api/types"

	"firewall_containers/network-go/config"
	"firewall_containers/network-go/docker"
	"firewall_containers/network-go/iptables"
)

// Compile-time interface conformance checks
var _ docker.DockerAPI = (*MockDockerClient)(nil)
var _ iptables.IPTablesAPI = (*MockIPTablesManager)(nil)

// MockDockerClient implements docker.DockerAPI for testing
type MockDockerClient struct {
	EnsureNetworkCalled           bool
	EnsureNetworkCfg              config.NetworkConfig
	EnsureNetworkErr              error

	ConnectContainerCalled        bool
	ConnectContainerName          string
	ConnectContainerNetwork       string
	ConnectContainerIP            string
	ConnectContainerErr           error

	WaitForRunningCalled          bool
	WaitForRunningName            string

	GetContainerPIDCalled         bool
	GetContainerPIDName           string
	GetContainerPIDResult         int
	GetContainerPIDErr            error

	AddRouteCalled                bool
	AddRouteContainer             string
	AddRouteNetwork               string
	AddRouteGateway               string
	AddRouteErr                   error

	FindContainerNameCalled       bool
	FindContainerNameResult       string
	FindContainerNameErr          error

	InspectContainerErr           error
	RemoveNetworkErr              error
	DisconnectContainerErr        error
}

func (m *MockDockerClient) Close() error { return nil }

func (m *MockDockerClient) EnsureNetwork(ctx context.Context, netCfg config.NetworkConfig) error {
	m.EnsureNetworkCalled = true
	m.EnsureNetworkCfg = netCfg
	return m.EnsureNetworkErr
}

func (m *MockDockerClient) RemoveNetwork(ctx context.Context, networkName string) error {
	return m.RemoveNetworkErr
}

func (m *MockDockerClient) ConnectContainer(ctx context.Context, containerName, networkName, ip string) error {
	m.ConnectContainerCalled = true
	m.ConnectContainerName = containerName
	m.ConnectContainerNetwork = networkName
	m.ConnectContainerIP = ip
	return m.ConnectContainerErr
}

func (m *MockDockerClient) DisconnectContainer(ctx context.Context, containerName, networkName string) error {
	return m.DisconnectContainerErr
}

func (m *MockDockerClient) InspectContainer(ctx context.Context, containerName string) (*types.ContainerJSON, error) {
	return nil, m.InspectContainerErr
}

func (m *MockDockerClient) WaitForContainerRunning(ctx context.Context, containerName string, timeout time.Duration) error {
	m.WaitForRunningCalled = true
	m.WaitForRunningName = containerName
	return nil
}

func (m *MockDockerClient) GetContainerPID(ctx context.Context, containerName string) (int, error) {
	m.GetContainerPIDCalled = true
	m.GetContainerPIDName = containerName
	return m.GetContainerPIDResult, m.GetContainerPIDErr
}

func (m *MockDockerClient) AddRouteInContainer(ctx context.Context, containerName, network, gateway string) error {
	m.AddRouteCalled = true
	m.AddRouteContainer = containerName
	m.AddRouteNetwork = network
	m.AddRouteGateway = gateway
	return m.AddRouteErr
}

func (m *MockDockerClient) FindContainerName(ctx context.Context, name, selector string) (string, error) {
	m.FindContainerNameCalled = true
	if m.FindContainerNameResult != "" {
		return m.FindContainerNameResult, m.FindContainerNameErr
	}
	return name, m.FindContainerNameErr
}

// MockIPTablesManager implements iptables.IPTablesAPI for testing
type MockIPTablesManager struct {
	BinaryResult                                string
	EnsureIPForwardCalled                       bool
	EnsureIPForwardErr                          error
	EnsureEstablishedRelatedCalled              bool
	EnsureEstablishedRelatedChain                string
	EnsureEstablishedRelatedErr                 error

	InsertPreroutingRuleCalled                  bool
	InsertPreroutingRuleArgs                    []string
	InsertPreroutingRuleErr                     error

	InsertPreroutingRuleOnInterfaceCalled       bool
	InsertPreroutingRuleOnInterfaceArgs         []string
	InsertPreroutingRuleOnInterfaceErr          error

	InsertPostroutingMasqueradeCalled           bool
	InsertPostroutingMasqueradeArgs             []string
	InsertPostroutingMasqueradeErr              error

	InsertForwardAcceptCalled                   bool
	InsertForwardAcceptChain                    string
	InsertForwardAcceptSourceIP                 string
	InsertForwardAcceptTargetIP                 string
	InsertForwardAcceptProto                    string
	InsertForwardAcceptSourcePort               string
	InsertForwardAcceptTargetPort               string
	InsertForwardAcceptComment                  string
	InsertForwardAcceptErr                      error

	InsertPreroutingRuleInContainerCalled       bool
	InsertPreroutingRuleInContainerPID          int
	InsertPreroutingRuleInContainerArgs         []string
	InsertPreroutingRuleInContainerErr          error

	InsertPostroutingMasqueradeInContainerCalled bool
	InsertPostroutingMasqueradeInContainerErr   error
	DeleteForwardAcceptErr                      error
	DeleteLineErr                               error
}

func (m *MockIPTablesManager) Binary() string {
	if m.BinaryResult == "" {
		return "/usr/sbin/iptables"
	}
	return m.BinaryResult
}

func (m *MockIPTablesManager) EnsureIPForward() error {
	m.EnsureIPForwardCalled = true
	return m.EnsureIPForwardErr
}

func (m *MockIPTablesManager) EnsureEstablishedRelated(chain string) error {
	m.EnsureEstablishedRelatedCalled = true
	m.EnsureEstablishedRelatedChain = chain
	return m.EnsureEstablishedRelatedErr
}

func (m *MockIPTablesManager) DeleteLine(chain string, lineNum string) error {
	return m.DeleteLineErr
}

func (m *MockIPTablesManager) InsertPreroutingRule(sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error {
	m.InsertPreroutingRuleCalled = true
	m.InsertPreroutingRuleArgs = []string{sourceIP, proto, sourcePort, targetIP, targetPort, comment}
	return m.InsertPreroutingRuleErr
}

func (m *MockIPTablesManager) InsertPreroutingRuleOnInterface(iface, proto, sourcePort, targetIP, targetPort, comment string) error {
	m.InsertPreroutingRuleOnInterfaceCalled = true
	m.InsertPreroutingRuleOnInterfaceArgs = []string{iface, proto, sourcePort, targetIP, targetPort, comment}
	return m.InsertPreroutingRuleOnInterfaceErr
}

func (m *MockIPTablesManager) InsertPostroutingMasquerade(sourceCIDR, proto, sourcePort, comment string) error {
	m.InsertPostroutingMasqueradeCalled = true
	m.InsertPostroutingMasqueradeArgs = []string{sourceCIDR, proto, sourcePort, comment}
	return m.InsertPostroutingMasqueradeErr
}

func (m *MockIPTablesManager) InsertPostroutingMasqueradeForTarget(targetCIDR, proto, targetPort, comment string) error {
	return nil
}

func (m *MockIPTablesManager) InsertForwardAccept(chain, sourceIP, targetIP, proto, sourcePort, targetPort, comment string) error {
	m.InsertForwardAcceptCalled = true
	m.InsertForwardAcceptChain = chain
	m.InsertForwardAcceptSourceIP = sourceIP
	m.InsertForwardAcceptTargetIP = targetIP
	m.InsertForwardAcceptProto = proto
	m.InsertForwardAcceptSourcePort = sourcePort
	m.InsertForwardAcceptTargetPort = targetPort
	m.InsertForwardAcceptComment = comment
	return m.InsertForwardAcceptErr
}

func (m *MockIPTablesManager) DeleteForwardAccept(chain, comment string) error {
	return m.DeleteForwardAcceptErr
}

func (m *MockIPTablesManager) InsertPreroutingRuleInContainer(pid int, sourceIP, proto, sourcePort, targetIP, targetPort, comment string) error {
	m.InsertPreroutingRuleInContainerCalled = true
	m.InsertPreroutingRuleInContainerPID = pid
	m.InsertPreroutingRuleInContainerArgs = []string{sourceIP, proto, sourcePort, targetIP, targetPort, comment}
	return m.InsertPreroutingRuleInContainerErr
}

func (m *MockIPTablesManager) InsertPostroutingMasqueradeInContainer(pid int, sourceCIDR, proto, sourcePort, comment string) error {
	m.InsertPostroutingMasqueradeInContainerCalled = true
	return m.InsertPostroutingMasqueradeInContainerErr
}