Repository: kubernetes/kubernetes
File: test/e2e/network/netpol/model.go (Go)

Code before:
// The *total* number of containers is namespaces x pods x ports x protocols.
func NewModel(namespaceNames []string, podNames []string, ports []int32, protocols []v1.Protocol) *Model {
	podNamesByNamespace := make(map[string][]string, len(namespaceNames))
	for _, ns := range namespaceNames {
		podNamesByNamespace[ns] = podNames
	}
	return newModelWithPerNamespacePodNames(namespaceNames, podNamesByNamespace, ports, protocols)
}

// newModelWithPerNamespacePodNames instantiates a model where each namespace can define a different pod set.
func newModelWithPerNamespacePodNames(namespaceNames []string, podNamesByNamespace map[string][]string, ports []int32, protocols []v1.Protocol) *Model {
	model := &Model{
		Ports:     ports,
		Protocols: protocols,
	}
	seenPodNames := map[string]struct{}{}

	// build the entire "model" for the overall test, which means, building
	// namespaces, pods, containers for each protocol.
	for _, ns := range namespaceNames {
		var pods []*Pod
		podNames := podNamesByNamespace[ns]
		for _, podName := range podNames {
			if _, seen := seenPodNames[podName]; !seen {
				model.PodNames = append(model.PodNames, podName)
				seenPodNames[podName] = struct{}{}
			}
			pods = append(pods, &Pod{
				Name:      podName,
				Ports:     ports,
				Protocols: protocols,
			})
		}
		model.Namespaces = append(model.Namespaces, &Namespace{
			Name: ns,
			Pods: pods,
		})
	}
	return model
}

// Namespace is the abstract representation of what matters to network policy
// tests for a namespace; i.e. it ignores kube implementation details
type Namespace struct {
	Name string
	Pods []*Pod
}

// Pod is the abstract representation of what matters to network policy tests for
// a pod; i.e. it ignores kube implementation details
type Pod struct {

Review comment (refactor):
if `podNamesByNamespace` is a `map[string]sets.Set[string]` then you don't need `seenPodNames` here, since it wouldn't be possible for it to contain duplicates.

Code after:
// The number of containers per pod is the number of ports x the number of protocols.
// The *total* number of containers is namespaces x pods x ports x protocols.
func NewModel(namespaceNames []string, podNames []string, ports []int32, protocols []v1.Protocol) *Model {
	podNamesByNamespace := make(map[string]sets.Set[string], len(namespaceNames))
	for _, ns := range namespaceNames {
		podNamesByNamespace[ns] = sets.New(podNames...)
	}
	return newModelWithPerNamespacePodNames(namespaceNames, podNamesByNamespace, ports, protocols)
}

// newModelWithPerNamespacePodNames instantiates a model where each namespace can define a different pod set.
func newModelWithPerNamespacePodNames(namespaceNames []string, podNamesByNamespace map[string]sets.Set[string], ports []int32, protocols []v1.Protocol) *Model {
	model := &Model{
		Ports:     ports,
		Protocols: protocols,
	}
	allPodNames := sets.New[string]()

	// build the entire "model" for the overall test, which means, building
	// namespaces, pods, containers for each protocol.
	for _, ns := range namespaceNames {
		var pods []*Pod
		podNames := sets.List(podNamesByNamespace[ns])
		for _, podName := range podNames {
			allPodNames.Insert(podName)
			pods = append(pods, &Pod{
				Name:      podName,
				Ports:     ports,
				Protocols: protocols,
			})
		}
		model.Namespaces = append(model.Namespaces, &Namespace{
			Name: ns,
			Pods: pods,
		})
	}
	model.PodNames = sets.List(allPodNames)
	return model
}

// Namespace is the abstract representation of what matters to network policy
// tests for a namespace; i.e. it ignores kube implementation details
type Namespace struct {
	Name string
	Pods []*Pod
}

// Pod is the abstract representation of what matters to network policy tests for
// a pod; i.e. it ignores kube implementation details
type Pod struct {
	Name      string