// Copyright (c) 2015-2020 Tigera, Inc. All rights reserved.

package main_test

import (
	"context"
	"fmt"
	"os"

	"github.com/containernetworking/cni/pkg/types"
	cniv1 "github.com/containernetworking/cni/pkg/types/100"
	"github.com/google/uuid"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	apiv3 "github.com/projectcalico/api/pkg/apis/projectcalico/v3"
	crclient "sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/projectcalico/calico/cni-plugin/internal/pkg/testutils"
	apiconfig "github.com/projectcalico/calico/libcalico-go/lib/apiconfig"
	"github.com/projectcalico/calico/libcalico-go/lib/apis/internalapi"
	"github.com/projectcalico/calico/libcalico-go/lib/backend/k8s"
	"github.com/projectcalico/calico/libcalico-go/lib/backend/k8s/rawcrdclient"
	client "github.com/projectcalico/calico/libcalico-go/lib/clientv3"
	"github.com/projectcalico/calico/libcalico-go/lib/ipam"
	"github.com/projectcalico/calico/libcalico-go/lib/ipam/ipamtestutils"
	"github.com/projectcalico/calico/libcalico-go/lib/names"
	cnet "github.com/projectcalico/calico/libcalico-go/lib/net"
	"github.com/projectcalico/calico/libcalico-go/lib/options"
)

var (
	plugin          = "calico-ipam"
	defaultIPv4Pool = "192.168.0.0/16"
)

// verifyRoutesPopulatedInResult verifies that routes are correctly populated in the IPAM result.
// For normal pods (expectNonEmpty=true), it checks that routes exist for all assigned IPs.
// For migration target pods (expectNonEmpty=false), it checks that routes are empty.
func verifyRoutesPopulatedInResult(result *cniv1.Result, expectNonEmpty bool) {
	if expectNonEmpty {
		// Normal pods should have routes for all assigned IPs
		Expect(result.Routes).NotTo(BeEmpty(), "Routes should be present for assigned IPs")

		// Verify that each assigned IP has a corresponding route whose destination
		// network contains the IP. We match by IP containment rather than exact string
		// equality because older CNI spec versions (< 0.3.0) normalise the IP mask to
		// /32 (or /128) during result conversion while routes keep the block-level mask.
		for _, ip := range result.IPs {
			found := false
			for _, route := range result.Routes {
				if route.Dst.Contains(ip.Address.IP) {
					found = true
					break
				}
			}
			Expect(found).To(BeTrue(),
				fmt.Sprintf("Route should exist for IP %s", ip.Address.String()))
		}
	} else {
		// Migration target pods should have empty routes
		Expect(result.Routes).To(BeEmpty(), "Routes should be empty for migration target pod")
	}
}

var _ = Describe("Calico IPAM Tests", func() {
	cniVersion := os.Getenv("CNI_SPEC_VERSION")
	config, err := apiconfig.LoadClientConfigFromEnvironment()
	Expect(err).NotTo(HaveOccurred())
	calicoClient, err := client.New(*config)
	Expect(err).NotTo(HaveOccurred())
	k8sClient := getKubernetesClient()

	var cid string
	BeforeEach(func() {
		testutils.WipeDatastore()
		testutils.MustCreateNewIPPool(calicoClient, defaultIPv4Pool, false, false, true)
		testutils.MustCreateNewIPPool(calicoClient, "fd80:24e2:f998:72d6::/64", false, false, true)

		// Create a unique container ID for each test.
		cid = uuid.NewString()

		// Create the node for these tests. The IPAM code requires a corresponding Calico node to exist.
		var name string
		name, err = names.Hostname()
		Expect(err).NotTo(HaveOccurred())
		err = testutils.AddNode(calicoClient, k8sClient, name)
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(func() {
		// Delete the node.
		name, err := names.Hostname()
		Expect(err).NotTo(HaveOccurred())
		err = testutils.DeleteNode(calicoClient, k8sClient, name)
		Expect(err).NotTo(HaveOccurred())
	})

	DescribeTable("Request different numbers of IP addresses",
		func(expectedIPv4, expectedIPv6 bool, netconf string) {
			result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
			var ip4Mask, ip6Mask string

			for _, ip := range result.IPs {
				if ip.Address.IP.To4() != nil {
					ip4Mask = ip.Address.Mask.String()
				} else if ip.Address.IP.To16() != nil {
					ip6Mask = ip.Address.Mask.String()
				}
			}

			if expectedIPv4 {
				Expect(ip4Mask).Should(Equal("ffffffc0"))
			}

			if expectedIPv6 {
				Expect(ip6Mask).Should(Equal("ffffffffffffffffffffffffffffffc0"))
			}

			// Verify that routes are present for each assigned IP
			verifyRoutesPopulatedInResult(result, true)

			_, _, exitCode := testutils.RunIPAMPlugin(netconf, "DEL", "", cid, cniVersion)
			Expect(exitCode).Should(Equal(0))
		},

		Entry("IPAM with no configuration", true, false, fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                  "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "log_level": "debug",
              "datastore_type": "%s",
              "ipam": {
                "type": "%s"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))),
		Entry("IPAM with IPv4 (explicit)", true, false, fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                 "kubeconfig": "/home/user/certs/kubeconfig"
                  },
              "datastore_type": "%s",
              "ipam": {
                "type": "%s",
                "assign_ipv4": "true"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))),
		Entry("IPAM with IPv6 only", false, true, fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                 "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "datastore_type": "%s",
              "ipam": {
                "type": "%s",
                "assign_ipv4": "false",
                "assign_ipv6": "true"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))),
		Entry("IPAM with IPv4 and IPv6", true, true, fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                 "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "datastore_type": "%s",
              "log_level": "debug",
              "ipam": {
                "type": "%s",
                "assign_ipv4": "true",
                "assign_ipv6": "true"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec)),
		),
	)

	Describe("Run IPAM plugin - No partial assignments in dual stack", func() {
		Context("With no IPv4 pool", func() {
			It("Should not assign an IPv6 address in a dual stack configuration", func() {
				// Delete IPv6 pool.
				testutils.MustDeleteIPPool(calicoClient, "fd80:24e2:f998:72d6::/64")

				// Assign an IP to fill up the pool
				netconf := fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                 "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "datastore_type": "%s",
              "log_level": "debug",
              "ipam": {
                "type": "%s",
                "assign_ipv4": "true",
                "assign_ipv6": "true"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				_, _, _ = testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)

				// Attempt to assign both an IPv4 and IPv6 address
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(result.IPs).To(HaveLen(0))

				// Clean up
				_, _, exitCode := testutils.RunIPAMPlugin(netconf, "DEL", "", cid, cniVersion)
				Expect(exitCode).Should(Equal(0))
				testutils.MustCreateNewIPPool(calicoClient, "fd80:24e2:f998:72d6::/64", false, false, true)
			})
		})

		Context("With no IPv6 pool", func() {
			It("Should not assign an IPv4 address in a dual stack configuration", func() {
				// Delete IPv4 pool.
				testutils.MustDeleteIPPool(calicoClient, defaultIPv4Pool)

				// Assign an IP to fill up the pool
				netconf := fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                 "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "datastore_type": "%s",
              "log_level": "debug",
              "ipam": {
                "type": "%s",
                "assign_ipv4": "true",
                "assign_ipv6": "true"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				_, _, _ = testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)

				// Attempt to assign both an IPv4 and IPv6 address
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(result.IPs).To(HaveLen(0))

				// Clean up
				_, _, exitCode := testutils.RunIPAMPlugin(netconf, "DEL", "", cid, cniVersion)
				Expect(exitCode).Should(Equal(0))
				testutils.MustCreateNewIPPool(calicoClient, defaultIPv4Pool, false, false, true)
			})
		})
	})

	Describe("Run IPAM plugin - Verify IP Pools", func() {
		Context("Pass valid pools", func() {
			It("Uses the ipv4 pool", func() {
				netconf := fmt.Sprintf(`
                {
                  "cniVersion": "%s",
                  "name": "net1",
                  "type": "calico",
                  "etcd_endpoints": "http://%s:2379",
                  "kubernetes": {
                    "kubeconfig": "/home/user/certs/kubeconfig"
                    },
                  "datastore_type": "%s",
                  "ipam": {
                    "type": "%s",
                    "assign_ipv4": "true",
                    "ipv4_pools": [ "192.168.0.0/16" ]
                  },
                  "calico_api_group": "%s"
                }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(len(result.IPs)).To(Equal(1))
				Expect(result.IPs[0].Address.IP.String()).Should(HavePrefix("192.168."))
			})
		})

		Context("Pass more than one pool", func() {
			It("Uses one of the ipv4 pools", func() {
				testutils.MustCreateNewIPPool(calicoClient, "192.169.1.0/24", false, false, true)
				netconf := fmt.Sprintf(`
                {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                        "kubernetes": {
                           "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s",
                        "assign_ipv4": "true",
                        "ipv4_pools": [ "192.169.1.0/24", "192.168.0.0/16" ]
                      },
                      "calico_api_group": "%s"
                }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(result.IPs[0].Address.IP.String()).Should(Or(HavePrefix("192.168."), HavePrefix("192.169.1")))
			})
		})

		Context("Disabled IP pool", func() {
			It("Never allocates from the disabled pool", func() {
				netconf := fmt.Sprintf(`
                {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                        "kubernetes": {
                         "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s",
                        "assign_ipv4": "true"
                      },
                      "calico_api_group": "%s"
                }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))

				// Get an allocation
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(result.IPs[0].Address.IP.String()).Should(HavePrefix("192.168."))

				// Disable the currently enabled pool
				pool, err := calicoClient.IPPools().Get(context.Background(), "192-168-0-0-16", options.GetOptions{})
				Expect(err).ToNot(HaveOccurred())
				pool.Spec.Disabled = true
				_, err = calicoClient.IPPools().Update(context.Background(), pool, options.SetOptions{})
				Expect(err).ToNot(HaveOccurred())

				// Create new (enabled) pool
				testutils.MustCreateNewIPPool(calicoClient, "192.169.1.0/24", false, false, true)

				// Get an allocation then check that it is not from the disabled pool
				result, _, _ = testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(result.IPs[0].Address.IP.String()).Should(HavePrefix("192.169.1"))

				// Re-enable the pool. We can't delete the node if the IP pool is disabled.
				// This is arguably a bug in the node deletion code...
				pool, err = calicoClient.IPPools().Get(context.Background(), "192-168-0-0-16", options.GetOptions{})
				Expect(err).ToNot(HaveOccurred())
				pool.Spec.Disabled = false
				_, err = calicoClient.IPPools().Update(context.Background(), pool, options.SetOptions{})
				Expect(err).ToNot(HaveOccurred())
			})
		})

		Context("Pass an invalid pool", func() {
			It("fails to get an IP", func() {
				// Put the bogus pool last in the array
				netconf := fmt.Sprintf(`
                    {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                        "kubernetes": {
                         "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s",
                        "assign_ipv4": "true",
                        "ipv4_pools": [ "192.168.0.0/16", "192.169.1.0/24" ]
                      },
                      "calico_api_group": "%s"
                    }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				_, err, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(err.Msg).Should(ContainSubstring("192.169.1.0/24) does not exist"))
			})

			It("fails to get an IP", func() {
				// Put the bogus pool first in the array
				netconf := fmt.Sprintf(`
                    {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                        "kubernetes": {
                         "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s",
                        "assign_ipv4": "true",
                        "ipv4_pools": [ "192.168.0.0/16", "192.169.1.0/24" ]
                      },
                      "calico_api_group": "%s"
                    }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
				_, err, _ := testutils.RunIPAMPlugin(netconf, "ADD", "", cid, cniVersion)
				Expect(err.Msg).Should(ContainSubstring("192.169.1.0/24) does not exist"))
			})
		})
	})

	Describe("Requesting an explicit IP address", func() {
		netconf := fmt.Sprintf(`
                    {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                      "kubernetes": {
                        "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s"
                      },
                      "calico_api_group": "%s"
                    }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
		Context("Pass explicit IP address", func() {
			It("Return the expected IP", func() {
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid, cniVersion)
				Expect(len(result.IPs)).Should(Equal(1))
				Expect(result.IPs[0].Address.String()).Should(Equal("192.168.123.123/32"))
			})
			It("Return the expected IP twice after deleting in the middle", func() {
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid, cniVersion)
				Expect(len(result.IPs)).Should(Equal(1))
				Expect(result.IPs[0].Address.String()).Should(Equal("192.168.123.123/32"))
				_, _, _ = testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.123", cid, cniVersion)
				result, _, _ = testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid, cniVersion)
				Expect(len(result.IPs)).Should(Equal(1))
				Expect(result.IPs[0].Address.String()).Should(Equal("192.168.123.123/32"))
			})
			It("Doesn't allow an explicit IP to be assigned twice", func() {
				result, _, _ := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid, cniVersion)
				Expect(len(result.IPs)).Should(Equal(1))
				Expect(result.IPs[0].Address.String()).Should(Equal("192.168.123.123/32"))
				_, _, exitCode := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid, cniVersion)
				Expect(exitCode).Should(BeNumerically(">", 0))
			})
		})
	})

	Describe("Run IPAM DEL", func() {
		netconf := fmt.Sprintf(`
                    {
                      "cniVersion": "%s",
                      "name": "net1",
                      "type": "calico",
                      "etcd_endpoints": "http://%s:2379",
                      "kubernetes": {
                        "kubeconfig": "/home/user/certs/kubeconfig"
                      },
                      "datastore_type": "%s",
                      "ipam": {
                        "type": "%s"
                      },
                      "calico_api_group": "%s"
                    }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))

		It("should exit successfully even if no address exists", func() {
			_, _, exitCode := testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.123", cid, cniVersion)
			Expect(exitCode).Should(Equal(0))
		})

		Context("when using old IPAM handle", func() {
			It("should remove the old handle", func() {
				// Create an IP using workload as the handle. This is the "old" way of
				// assigning IPs, prior to this PR:
				// https://github.com/projectcalico/cni-plugin/pull/446
				workload := cid
				assignArgs := ipam.AssignIPArgs{
					IP:       cnet.MustParseIP("192.168.123.123"),
					HandleID: &workload,
				}
				ctx := context.Background()
				err := calicoClient.IPAM().AssignIP(ctx, assignArgs)
				Expect(err).NotTo(HaveOccurred())

				// Verify the new IP was set.
				ips, err := calicoClient.IPAM().IPsByHandle(ctx, workload)
				Expect(err).NotTo(HaveOccurred())
				Expect(ips).To(HaveLen(1))

				// Call the IPAM plugin and assert that it properly cleans up the allocation.
				result, e, rc := testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.123", cid, cniVersion)
				Expect(e).To(Equal(types.Error{}))
				Expect(rc).To(Equal(0))
				Expect(len(result.IPs)).Should(Equal(0))

				// Verify that the workload handle is gone.
				_, err = calicoClient.IPAM().IPsByHandle(ctx, workload)
				Expect(err).To(HaveOccurred())

				// Create an IP using the "new" style handleID composed of the
				// network name and containerID.
				handleID := fmt.Sprintf("net1.%s", cid)
				assignArgs = ipam.AssignIPArgs{
					IP:       cnet.MustParseIP("192.168.123.123"),
					HandleID: &handleID,
				}
				err = calicoClient.IPAM().AssignIP(ctx, assignArgs)
				Expect(err).NotTo(HaveOccurred())

				// Verify the new IP was set.
				ips, err = calicoClient.IPAM().IPsByHandle(ctx, handleID)
				Expect(err).NotTo(HaveOccurred())
				Expect(ips).To(HaveLen(1))

				// Remove the IP and handle.
				result, e, rc = testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.123", cid, cniVersion)
				Expect(e).To(Equal(types.Error{}))
				Expect(rc).To(Equal(0))
				Expect(len(result.IPs)).Should(Equal(0))

				// Verify that the handleID is gone.
				_, err = calicoClient.IPAM().IPsByHandle(ctx, handleID)
				Expect(err).To(HaveOccurred())
			})
		})
	})

	Describe("Upgrade of block affinities occurs only once and creates touchfile", func() {
		var (
			crdClient crclient.Client
			host      string
			netconf   string
			useV3     bool
		)

		BeforeEach(func() {
			if os.Getenv("DATASTORE_TYPE") != "kubernetes" {
				Skip("Upgrade behavior is only relevant on Kubernetes datastore")
			}

			// Ensure clean environment: remove touchfile if it exists.
			_ = os.Remove("/var/run/calico/cni/ipam_upgraded")

			// Build a raw CRD REST client using the same kubeconfig as other tests.
			cfg, err := apiconfig.LoadClientConfigFromEnvironment()
			Expect(err).NotTo(HaveOccurred())
			kcfg, _, err := k8s.CreateKubernetesClientset(&cfg.Spec)
			Expect(err).NotTo(HaveOccurred())
			useV3 = k8s.BackendAPIGroup(&config.Spec) == "projectcalico.org/v3"
			crdClient, err = rawcrdclient.New(kcfg, useV3)
			Expect(err).NotTo(HaveOccurred())

			// Determine host for BlockAffinity and ensure there is a usable netconf.
			host, err = names.Hostname()
			Expect(err).NotTo(HaveOccurred())

			cniVersion := os.Getenv("CNI_SPEC_VERSION")
			netconf = fmt.Sprintf(`
            {
              "cniVersion": "%s",
              "name": "net1",
              "type": "calico",
              "etcd_endpoints": "http://%s:2379",
              "kubernetes": {
                "kubeconfig": "/home/user/certs/kubeconfig"
              },
              "datastore_type": "%s",
              "log_level": "debug",
              "ipam": {
                "type": "%s"
              },
              "calico_api_group": "%s"
            }`, cniVersion, os.Getenv("ETCD_IP"), os.Getenv("DATASTORE_TYPE"), plugin, k8s.BackendAPIGroup(&config.Spec))
		})

		It("upgrades unlabeled block affinities on first call and not on subsequent calls", func() {
			ctx := context.Background()
			// 1) Create an unlabeled BlockAffinity for this host.
			Expect(ipamtestutils.CreateUnlabeledBlockAffinity(ctx, crdClient, host, "10.11.0.0/26")).To(Succeed())

			getLabels := func(cidr string) map[string]string {
				if useV3 {
					var list apiv3.BlockAffinityList
					Expect(crdClient.List(ctx, &list)).To(Succeed())
					var found1 *apiv3.BlockAffinity
					for i := range list.Items {
						if list.Items[i].Spec.Node == host && list.Items[i].Spec.CIDR == cidr {
							found1 = &list.Items[i]
							break
						}
					}
					ExpectWithOffset(1, found1).NotTo(BeNil())
					return found1.Labels
				} else {
					var list internalapi.BlockAffinityList
					Expect(crdClient.List(ctx, &list)).To(Succeed())
					var found1 *internalapi.BlockAffinity
					for i := range list.Items {
						if list.Items[i].Spec.Node == host && list.Items[i].Spec.CIDR == cidr {
							found1 = &list.Items[i]
							break
						}
					}
					ExpectWithOffset(1, found1).NotTo(BeNil())
					return found1.Labels
				}
			}

			// Sanity check: ensure the BA has no labels.
			foundLabels := getLabels("10.11.0.0/26")
			Expect(foundLabels).To(HaveLen(0))

			// 2) Run the IPAM plugin with explicit IP to trigger maybeUpgradeIPAM.
			cid1 := uuid.NewString()
			result, errOut, exit := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.123", cid1, os.Getenv("CNI_SPEC_VERSION"))
			Expect(exit).To(Equal(0), fmt.Sprintf("stderr: %+v", errOut))
			Expect(result.IPs).To(HaveLen(1))

			// 3) Verify the touchfile exists.
			_, statErr := os.Stat("/var/run/calico/cni/ipam_upgraded")
			Expect(statErr).NotTo(HaveOccurred())

			// 4) Verify the previously unlabeled BA is now labeled.
			foundLabels = getLabels("10.11.0.0/26")
			Expect(foundLabels).To(HaveKey(apiv3.LabelAffinityType))
			Expect(foundLabels).To(HaveKey(apiv3.LabelIPVersion))
			Expect(foundLabels).To(HaveKey(apiv3.LabelHostnameHash))

			// 5) Create a second unlabeled BA for this host.
			Expect(ipamtestutils.CreateUnlabeledBlockAffinity(ctx, crdClient, host, "10.11.0.64/26")).To(Succeed())

			labels2 := getLabels("10.11.0.64/26")
			Expect(labels2).To(HaveLen(0))

			// 6) Run the plugin again; since touchfile exists, no upgrade should occur.
			cid2 := uuid.NewString()
			result2, errOut2, exit2 := testutils.RunIPAMPlugin(netconf, "ADD", "IP=192.168.123.124", cid2, os.Getenv("CNI_SPEC_VERSION"))
			Expect(exit2).To(Equal(0), fmt.Sprintf("stderr: %+v", errOut2))
			Expect(result2.IPs).To(HaveLen(1))

			// Re-list and ensure the second BA remains unlabeled.
			labels2 = getLabels("10.11.0.64/26")
			Expect(labels2).To(HaveLen(0))

			// Cleanup: release explicit IPs.
			_, _, exit = testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.123", cid1, os.Getenv("CNI_SPEC_VERSION"))
			Expect(exit).To(Equal(0))
			_, _, exit2 = testutils.RunIPAMPlugin(netconf, "DEL", "IP=192.168.123.124", cid2, os.Getenv("CNI_SPEC_VERSION"))
			Expect(exit2).To(Equal(0))
		})
	})
})
