towerops/test/towerops/agents_test.exs

defmodule Towerops.AgentsTest do
  use Towerops.DataCase

  import Towerops.AccountsFixtures

  alias Towerops.Agents
  alias Towerops.Agents.AgentAssignment
  alias Towerops.Agents.AgentToken

  setup do
    user = user_fixture()

    {:ok, organization} =
      Towerops.Organizations.create_organization(%{name: "Test Org"}, user.id)

    %{organization: organization, user: user}
  end

  describe "create_agent_token/2" do
    test "creates agent token with valid organization and name", %{organization: org} do
      assert {:ok, agent_token, token} = Agents.create_agent_token(org.id, "Test Agent")
      assert is_binary(token)
      assert byte_size(token) > 0
      assert agent_token.organization_id == org.id
      assert agent_token.name == "Test Agent"
      assert agent_token.enabled == true
      assert is_binary(agent_token.token)
      assert agent_token.token == token
    end

    test "returns error with invalid data" do
      assert {:error, changeset} = Agents.create_agent_token(nil, "")
      refute changeset.valid?
    end
  end

  describe "list_organization_agent_tokens/1" do
    test "returns all agent tokens for organization", %{organization: org1, user: user} do
      {:ok, org2} = Towerops.Organizations.create_organization(%{name: "Test Org 2"}, user.id)

      {:ok, token1, _} = Agents.create_agent_token(org1.id, "Agent 1")
      {:ok, token2, _} = Agents.create_agent_token(org1.id, "Agent 2")
      {:ok, _token3, _} = Agents.create_agent_token(org2.id, "Agent 3")

      tokens = Agents.list_organization_agent_tokens(org1.id)

      assert length(tokens) == 2
      assert Enum.any?(tokens, &(&1.id == token1.id))
      assert Enum.any?(tokens, &(&1.id == token2.id))
    end

    test "returns empty list for organization with no tokens", %{organization: org} do
      assert Agents.list_organization_agent_tokens(org.id) == []
    end
  end

  describe "verify_agent_token/1" do
    test "verifies valid token", %{organization: org} do
      {:ok, agent_token, token} = Agents.create_agent_token(org.id, "Test Agent")

      assert {:ok, verified_token} = Agents.verify_agent_token(token)
      assert verified_token.id == agent_token.id
    end

    test "fails for invalid token" do
      assert {:error, :invalid_token} = Agents.verify_agent_token("invalid_token")
    end

    test "fails for revoked token", %{organization: org} do
      {:ok, agent_token, token} = Agents.create_agent_token(org.id, "Test Agent")
      {:ok, _} = Agents.revoke_agent_token(agent_token.id)

      assert {:error, :invalid_token} = Agents.verify_agent_token(token)
    end
  end

  describe "revoke_agent_token/1" do
    test "disables agent token", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      assert {:ok, revoked_token} = Agents.revoke_agent_token(agent_token.id)
      assert revoked_token.enabled == false
    end
  end

  describe "update_agent_token/2" do
    test "updates agent token name", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Old Name")

      assert {:ok, updated_token} = Agents.update_agent_token(agent_token, %{name: "New Name"})
      assert updated_token.name == "New Name"
      assert updated_token.id == agent_token.id
    end

    test "returns error with invalid data", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Name")

      assert {:error, changeset} = Agents.update_agent_token(agent_token, %{name: ""})
      refute changeset.valid?
      assert "can't be blank" in errors_on(changeset).name
    end

    test "does not update other fields", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Name")
      original_token = agent_token.token

      {:ok, updated_token} =
        Agents.update_agent_token(agent_token, %{name: "Updated Name", token: "hacked"})

      assert updated_token.name == "Updated Name"
      assert updated_token.token == original_token
    end
  end

  describe "delete_agent_token/1" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      %{site: site, agent_token: agent_token}
    end

    test "deletes agent token from database", %{agent_token: agent_token} do
      assert {:ok, _} = Agents.delete_agent_token(agent_token.id)

      assert_raise Ecto.NoResultsError, fn ->
        Agents.get_agent_token!(agent_token.id)
      end
    end

    test "removes all direct device assignments", %{
      site: site,
      agent_token: agent_token
    } do
      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Device 1",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Device 2",
          ip_address: "192.168.1.2",
          site_id: site.id
        })

      # Assign both devices to the agent
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device1.id)
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device2.id)

      assert Agents.count_assigned_devices(agent_token.id) == 2

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token.id)

      # Verify assignments are removed
      assert Agents.get_device_assignment(device1.id) == nil
      assert Agents.get_device_assignment(device2.id) == nil
    end

    test "devices fall back to site agent when available", %{
      organization: org,
      site: site,
      agent_token: agent_token_to_delete
    } do
      # Create another agent that will be the site default
      {:ok, site_agent, _} = Agents.create_agent_token(org.id, "Site Agent")

      # Set site to use the site_agent
      {:ok, _site} = Towerops.Sites.update_site(site, %{agent_token_id: site_agent.id})

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Device",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      # Directly assign device to agent_to_delete (overriding site default)
      {:ok, _} = Agents.assign_device_to_agent(agent_token_to_delete.id, device.id)

      # Device should currently use agent_to_delete
      device = Repo.preload(device, site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == agent_token_to_delete.id

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token_to_delete.id)

      # Device should now fall back to site agent
      device = device |> Repo.reload!() |> Repo.preload(site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == site_agent.id
    end

    test "devices fall back to organization default when available", %{
      organization: org,
      site: site,
      agent_token: agent_token_to_delete
    } do
      # Create another agent that will be the org default
      {:ok, org_agent, _} = Agents.create_agent_token(org.id, "Org Agent")

      # Set organization to use the org_agent
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: org_agent.id})

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Device",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      # Directly assign device to agent_to_delete (overriding org default)
      {:ok, _} = Agents.assign_device_to_agent(agent_token_to_delete.id, device.id)

      # Device should currently use agent_to_delete
      device = Repo.preload(device, site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == agent_token_to_delete.id

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token_to_delete.id)

      # Device should now fall back to org default
      device = device |> Repo.reload!() |> Repo.preload(site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == org_agent.id
    end

    test "devices have no agent when no fallback exists (cloud polling)", %{
      site: site,
      agent_token: agent_token
    } do
      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Device",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      # Directly assign device to agent (no site or org defaults)
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device.id)

      # Device should currently use the agent
      device = Repo.preload(device, site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == agent_token.id

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token.id)

      # Device should now have no agent (cloud polling)
      device = device |> Repo.reload!() |> Repo.preload(site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == nil
    end

    test "removes agent from site defaults when deleting", %{
      organization: org,
      site: site,
      agent_token: agent_token
    } do
      # Set site to use this agent
      {:ok, _site} = Towerops.Sites.update_site(site, %{agent_token_id: agent_token.id})

      # Create another agent as fallback
      {:ok, fallback_agent, _} = Agents.create_agent_token(org.id, "Fallback Agent")

      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: fallback_agent.id})

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Device",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      # Device should use site agent
      device = Repo.preload(device, site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == agent_token.id

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token.id)

      # Site should no longer have this agent
      site = Repo.reload!(site)
      assert site.agent_token_id == nil

      # Device should fall back to org default
      device = device |> Repo.reload!() |> Repo.preload(site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == fallback_agent.id
    end

    test "removes agent from organization defaults when deleting", %{
      organization: org,
      site: site,
      agent_token: agent_token
    } do
      # Set organization to use this agent
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent_token.id})

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Device",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      # Device should use org default
      device = Repo.preload(device, site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == agent_token.id

      # Delete the agent
      assert {:ok, _} = Agents.delete_agent_token(agent_token.id)

      # Organization should no longer have this agent as default
      org = Repo.reload!(org)
      assert org.default_agent_token_id == nil

      # Device should have no agent
      device = device |> Repo.reload!() |> Repo.preload(site: [organization: :default_agent_token])
      assert Agents.get_effective_agent_token(device) == nil
    end
  end

  describe "update_agent_token_heartbeat/3" do
    test "updates last_seen_at and metadata", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      metadata = %{version: "0.1.0", hostname: "test-host"}
      {1, _} = Agents.update_agent_token_heartbeat(agent_token.id, "192.168.1.1", metadata)

      updated_token = Repo.get!(AgentToken, agent_token.id)
      assert updated_token.last_ip == "192.168.1.1"
      assert updated_token.metadata["version"] == "0.1.0"
      assert updated_token.metadata["hostname"] == "test-host"
      assert updated_token.last_seen_at
    end
  end

  describe "assign_device_to_agent/2" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      %{site: site, device: device}
    end

    test "assigns equipment to agent", %{organization: org, device: device} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      assert {:ok, assignment} =
               Agents.assign_device_to_agent(agent_token.id, device.id)

      assert assignment.agent_token_id == agent_token.id
      assert assignment.device_id == device.id
      assert assignment.enabled == true
    end

    test "prevents assigning equipment to multiple agents", %{
      organization: org,
      device: device
    } do
      {:ok, agent1, _} = Agents.create_agent_token(org.id, "Agent 1")
      {:ok, agent2, _} = Agents.create_agent_token(org.id, "Agent 2")

      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device.id)

      assert {:error, changeset} = Agents.assign_device_to_agent(agent2.id, device.id)
      assert "has already been taken" in errors_on(changeset).device_id
    end
  end

  describe "unassign_device/1" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      %{site: site, device: device, agent_token: agent_token}
    end

    test "removes equipment assignment", %{agent_token: agent_token, device: device} do
      {:ok, _assignment} = Agents.assign_device_to_agent(agent_token.id, device.id)
      {1, _} = Agents.unassign_device(device.id)

      assert Repo.get_by(AgentAssignment, device_id: device.id) == nil
    end
  end

  describe "get_agent_token!/1" do
    test "returns agent token by id", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      retrieved_token = Agents.get_agent_token!(agent_token.id)
      assert retrieved_token.id == agent_token.id
      assert retrieved_token.name == "Test Agent"
    end

    test "raises when token not found" do
      assert_raise Ecto.NoResultsError, fn ->
        Agents.get_agent_token!(Ecto.UUID.generate())
      end
    end
  end

  describe "update_agent_token_heartbeat/3 edge cases" do
    test "updates without metadata when nil is provided", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      {1, _} = Agents.update_agent_token_heartbeat(agent_token.id, "192.168.1.1", nil)

      updated_token = Repo.get!(AgentToken, agent_token.id)
      assert updated_token.last_ip == "192.168.1.1"
      assert updated_token.last_seen_at
    end

    test "updates without metadata when empty map is provided", %{organization: org} do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      {1, _} = Agents.update_agent_token_heartbeat(agent_token.id, "192.168.1.1", %{})

      updated_token = Repo.get!(AgentToken, agent_token.id)
      assert updated_token.last_ip == "192.168.1.1"
      assert updated_token.last_seen_at
    end
  end

  describe "get_device_assignment/1" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      %{site: site, device: device, agent_token: agent_token}
    end

    test "returns assignment when equipment is assigned", %{
      agent_token: agent_token,
      device: device
    } do
      {:ok, assignment} = Agents.assign_device_to_agent(agent_token.id, device.id)

      retrieved_assignment = Agents.get_device_assignment(device.id)
      assert retrieved_assignment.id == assignment.id
    end

    test "returns nil when equipment is not assigned", %{device: device} do
      assert Agents.get_device_assignment(device.id) == nil
    end
  end

  describe "update_device_assignment/2" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, agent_token1, _} = Agents.create_agent_token(org.id, "Agent 1")
      {:ok, agent_token2, _} = Agents.create_agent_token(org.id, "Agent 2")

      %{
        site: site,
        device: device,
        agent_token1: agent_token1,
        agent_token2: agent_token2
      }
    end

    test "creates new assignment when equipment is unassigned", %{
      device: device,
      agent_token1: agent_token1
    } do
      assert {:ok, assignment} =
               Agents.update_device_assignment(device.id, agent_token1.id)

      assert assignment.agent_token_id == agent_token1.id
      assert assignment.device_id == device.id
    end

    test "updates existing assignment to new agent", %{
      device: device,
      agent_token1: agent_token1,
      agent_token2: agent_token2
    } do
      {:ok, _} = Agents.assign_device_to_agent(agent_token1.id, device.id)

      assert {:ok, updated_assignment} =
               Agents.update_device_assignment(device.id, agent_token2.id)

      assert updated_assignment.agent_token_id == agent_token2.id
      assert updated_assignment.device_id == device.id
    end

    test "removes assignment when set to nil", %{
      device: device,
      agent_token1: agent_token1
    } do
      {:ok, _} = Agents.assign_device_to_agent(agent_token1.id, device.id)

      assert {:ok, nil} = Agents.update_device_assignment(device.id, nil)
      assert Agents.get_device_assignment(device.id) == nil
    end
  end

  describe "list_agent_devices/1" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site.id
        })

      {:ok, device3} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 3",
          ip_address: "192.168.1.3",
          site_id: site.id
        })

      %{site: site, device1: device1, device2: device2, device3: device3}
    end

    test "returns devices assigned to agent with preloads", %{
      organization: org,
      device1: device1,
      device2: device2,
      device3: device3
    } do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device1.id)
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device2.id)

      device_list = Agents.list_agent_devices(agent_token.id)

      assert length(device_list) == 2
      assert Enum.any?(device_list, &(&1.id == device1.id))
      assert Enum.any?(device_list, &(&1.id == device2.id))
      refute Enum.any?(device_list, &(&1.id == device3.id))

      # Check preloads
      device = hd(device_list)
      assert Ecto.assoc_loaded?(device.site)
    end
  end

  describe "get_effective_agent_token/1 - hierarchical resolution" do
    setup %{organization: org} do
      {:ok, agent1, _} = Agents.create_agent_token(org.id, "Agent 1")
      {:ok, agent2, _} = Agents.create_agent_token(org.id, "Agent 2")
      {:ok, agent3, _} = Agents.create_agent_token(org.id, "Agent 3")

      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      %{
        agent1: agent1,
        agent2: agent2,
        agent3: agent3,
        site: site,
        device: device
      }
    end

    test "returns devices-level assignment (highest priority)", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      agent3: agent3,
      site: site,
      device: device
    } do
      # Set up hierarchy: org default, site override, equipment override
      {:ok, _org} = Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})
      {:ok, _site} = Towerops.Sites.update_site(site, %{agent_token_id: agent2.id})
      {:ok, _} = Agents.assign_device_to_agent(agent3.id, device.id)

      # Preload necessary associations
      device = Repo.preload(device, site: [organization: :default_agent_token])

      # Equipment assignment should take precedence
      assert Agents.get_effective_agent_token(device) == agent3.id
    end

    test "returns site-level assignment when equipment not assigned", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      site: site,
      device: device
    } do
      # Set up hierarchy: org default and site override, but no equipment assignment
      {:ok, _org} = Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})
      {:ok, _site} = Towerops.Sites.update_site(site, %{agent_token_id: agent2.id})

      # Preload necessary associations
      device = Repo.preload(device, site: [organization: :default_agent_token])

      # Site assignment should be used
      assert Agents.get_effective_agent_token(device) == agent2.id
    end

    test "returns organization-level assignment when site and equipment not assigned", %{
      organization: org,
      agent1: agent1,
      device: device
    } do
      # Set up hierarchy: only org default
      {:ok, _org} = Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      # Preload necessary associations
      device = Repo.preload(device, site: [organization: :default_agent_token])

      # Organization default should be used
      assert Agents.get_effective_agent_token(device) == agent1.id
    end

    test "returns nil when no agent assigned at any level", %{device: device} do
      # No assignments at any level
      device = Repo.preload(device, site: [organization: :default_agent_token])

      assert Agents.get_effective_agent_token(device) == nil
    end
  end

  describe "get_effective_agent_token_with_source/1" do
    setup %{organization: org} do
      {:ok, agent1, _} = Agents.create_agent_token(org.id, "Agent 1")
      {:ok, agent2, _} = Agents.create_agent_token(org.id, "Agent 2")
      {:ok, agent3, _} = Agents.create_agent_token(org.id, "Agent 3")

      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      %{
        agent1: agent1,
        agent2: agent2,
        agent3: agent3,
        site: site,
        device: device
      }
    end

    test "returns devices source for device-level assignment", %{
      agent3: agent3,
      device: device
    } do
      {:ok, _} = Agents.assign_device_to_agent(agent3.id, device.id)

      device = Repo.preload(device, site: [organization: :default_agent_token])

      assert {agent_id, :device} = Agents.get_effective_agent_token_with_source(device)
      assert agent_id == agent3.id
    end

    test "returns site source for site-level assignment", %{
      agent2: agent2,
      site: site,
      device: device
    } do
      {:ok, _site} = Towerops.Sites.update_site(site, %{agent_token_id: agent2.id})

      device = Repo.preload(device, site: [organization: :default_agent_token])

      assert {agent_id, :site} = Agents.get_effective_agent_token_with_source(device)
      assert agent_id == agent2.id
    end

    test "returns organization source for organization-level assignment", %{
      organization: org,
      agent1: agent1,
      device: device
    } do
      {:ok, _org} = Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      device = Repo.preload(device, site: [organization: :default_agent_token])

      assert {agent_id, :organization} = Agents.get_effective_agent_token_with_source(device)
      assert agent_id == agent1.id
    end

    test "returns none source when no assignment", %{device: device} do
      device = Repo.preload(device, site: [organization: :default_agent_token])

      assert {nil, :none} = Agents.get_effective_agent_token_with_source(device)
    end
  end

  describe "count_assigned_devices/1" do
    setup %{organization: org} do
      {:ok, site} =
        Towerops.Sites.create_site(%{
          name: "Test Site",
          organization_id: org.id
        })

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site.id
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site.id
        })

      {:ok, device3} =
        Towerops.Devices.create_device(%{
          name: "Test Equipment 3",
          ip_address: "192.168.1.3",
          site_id: site.id
        })

      %{site: site, device1: device1, device2: device2, device3: device3}
    end

    test "counts directly assigned equipment", %{
      organization: org,
      device1: device1,
      device2: device2
    } do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")

      # Initially no assignments
      assert Agents.count_assigned_devices(agent_token.id) == 0

      # Assign two pieces of equipment
      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device1.id)
      assert Agents.count_assigned_devices(agent_token.id) == 1

      {:ok, _} = Agents.assign_device_to_agent(agent_token.id, device2.id)
      assert Agents.count_assigned_devices(agent_token.id) == 2
    end

    test "does not count disabled assignments", %{
      organization: org,
      device1: device1
    } do
      {:ok, agent_token, _token} = Agents.create_agent_token(org.id, "Test Agent")
      {:ok, assignment} = Agents.assign_device_to_agent(agent_token.id, device1.id)

      assert Agents.count_assigned_devices(agent_token.id) == 1

      # Disable the assignment
      assignment
      |> Ecto.Changeset.change(enabled: false)
      |> Repo.update!()

      assert Agents.count_assigned_devices(agent_token.id) == 0
    end
  end

  describe "list_agent_polling_targets/1 - hierarchical polling" do
    setup %{organization: org} do
      {:ok, agent1, _} = Agents.create_agent_token(org.id, "Agent 1")
      {:ok, agent2, _} = Agents.create_agent_token(org.id, "Agent 2")
      {:ok, agent3, _} = Agents.create_agent_token(org.id, "Agent 3")

      {:ok, site1} =
        Towerops.Sites.create_site(%{
          name: "Site 1",
          organization_id: org.id
        })

      {:ok, site2} =
        Towerops.Sites.create_site(%{
          name: "Site 2",
          organization_id: org.id
        })

      %{
        agent1: agent1,
        agent2: agent2,
        agent3: agent3,
        site1: site1,
        site2: site2
      }
    end

    test "returns device with direct agent assignment", %{
      agent1: agent1,
      agent2: agent2,
      site1: site1
    } do
      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Assign equipment to agent1
      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device1.id)
      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device2.id)

      # Agent1 should see both equipment
      targets = Agents.list_agent_polling_targets(agent1.id)
      assert length(targets) == 2
      assert Enum.any?(targets, &(&1.id == device1.id))
      assert Enum.any?(targets, &(&1.id == device2.id))

      # Agent2 should see nothing
      targets = Agents.list_agent_polling_targets(agent2.id)
      assert targets == []
    end

    test "returns device inheriting from site agent", %{
      agent1: agent1,
      agent2: agent2,
      site1: site1
    } do
      # Set site1 to use agent1
      {:ok, _site} = Towerops.Sites.update_site(site1, %{agent_token_id: agent1.id})

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Agent1 should see both equipment (inherited from site)
      targets = Agents.list_agent_polling_targets(agent1.id)
      assert length(targets) == 2
      assert Enum.any?(targets, &(&1.id == device1.id))
      assert Enum.any?(targets, &(&1.id == device2.id))

      # Agent2 should see nothing
      targets = Agents.list_agent_polling_targets(agent2.id)
      assert targets == []
    end

    test "returns device inheriting from organization default", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      site1: site1
    } do
      # Set organization default to agent1
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Agent1 should see both equipment (inherited from organization)
      targets = Agents.list_agent_polling_targets(agent1.id)
      assert length(targets) == 2
      assert Enum.any?(targets, &(&1.id == device1.id))
      assert Enum.any?(targets, &(&1.id == device2.id))

      # Agent2 should see nothing
      targets = Agents.list_agent_polling_targets(agent2.id)
      assert targets == []
    end

    test "equipment-level assignment overrides site and organization", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      agent3: agent3,
      site1: site1
    } do
      # Set org default to agent1, site to agent2, equipment to agent3
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      {:ok, _site} = Towerops.Sites.update_site(site1, %{agent_token_id: agent2.id})

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Directly assign to agent3
      {:ok, _} = Agents.assign_device_to_agent(agent3.id, device1.id)

      # Only agent3 should see the equipment
      targets = Agents.list_agent_polling_targets(agent3.id)
      assert length(targets) == 1
      assert hd(targets).id == device1.id

      # Agent1 and agent2 should not see it
      assert Agents.list_agent_polling_targets(agent1.id) == []
      assert Agents.list_agent_polling_targets(agent2.id) == []
    end

    test "site-level assignment overrides organization default", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      site1: site1
    } do
      # Set org default to agent1, site to agent2
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      {:ok, _site} = Towerops.Sites.update_site(site1, %{agent_token_id: agent2.id})

      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Agent2 should see the equipment (site level)
      targets = Agents.list_agent_polling_targets(agent2.id)
      assert length(targets) == 1
      assert hd(targets).id == device1.id

      # Agent1 should not see it (org default is overridden by site)
      assert Agents.list_agent_polling_targets(agent1.id) == []
    end

    test "only returns SNMP-enabled devices", %{
      agent1: agent1,
      site1: site1
    } do
      # Equipment with SNMP enabled
      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Equipment with SNMP disabled
      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site1.id,
          snmp_enabled: false
        })

      # Assign both to agent1
      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device1.id)
      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device2.id)

      # Agent1 should only see SNMP-enabled devices
      targets = Agents.list_agent_polling_targets(agent1.id)
      assert length(targets) == 1
      assert hd(targets).id == device1.id
      assert hd(targets).snmp_enabled == true
    end

    test "returns device from multiple sites with different assignment levels", %{
      organization: org,
      agent1: agent1,
      agent2: agent2,
      site1: site1,
      site2: site2
    } do
      # Set org default to agent1
      {:ok, _org} =
        Towerops.Organizations.update_organization(org, %{default_agent_token_id: agent1.id})

      # Site2 overrides to agent2
      {:ok, _site} = Towerops.Sites.update_site(site2, %{agent_token_id: agent2.id})

      # Equipment at site1 (inherits org default - agent1)
      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Equipment at site2 (inherits site default - agent2)
      {:ok, device2} =
        Towerops.Devices.create_device(%{
          name: "Equipment 2",
          ip_address: "192.168.1.2",
          site_id: site2.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      # Agent1 should see device1 (from site1)
      targets = Agents.list_agent_polling_targets(agent1.id)
      assert length(targets) == 1
      assert hd(targets).id == device1.id

      # Agent2 should see device2 (from site2)
      targets = Agents.list_agent_polling_targets(agent2.id)
      assert length(targets) == 1
      assert hd(targets).id == device2.id
    end

    test "preloads necessary associations for API response", %{
      agent1: agent1,
      site1: site1
    } do
      {:ok, device1} =
        Towerops.Devices.create_device(%{
          name: "Equipment 1",
          ip_address: "192.168.1.1",
          site_id: site1.id,
          snmp_enabled: true,
          snmp_version: "2c",
          snmp_community: "public"
        })

      {:ok, _} = Agents.assign_device_to_agent(agent1.id, device1.id)

      targets = Agents.list_agent_polling_targets(agent1.id)
      device = hd(targets)

      # Verify preloads
      assert Ecto.assoc_loaded?(device.site)
      assert Ecto.assoc_loaded?(device.site.organization)
    end
  end
end