diff --git a/test/towerops/agents_test.exs b/test/towerops/agents_test.exs
index 7c815077..64675d47 100644
--- a/test/towerops/agents_test.exs
+++ b/test/towerops/agents_test.exs
@@ -168,6 +168,141 @@ defmodule Towerops.AgentsTest do
     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_equipment_assignment/1" do
+    setup %{organization: org} do
+      {:ok, site} =
+        Towerops.Sites.create_site(%{
+          name: "Test Site",
+          organization_id: org.id
+        })
+
+      {:ok, equipment} =
+        Towerops.Equipment.create_equipment(%{
+          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, equipment: equipment, agent_token: agent_token}
+    end
+
+    test "returns assignment when equipment is assigned", %{
+      agent_token: agent_token,
+      equipment: equipment
+    } do
+      {:ok, assignment} = Agents.assign_equipment_to_agent(agent_token.id, equipment.id)
+
+      retrieved_assignment = Agents.get_equipment_assignment(equipment.id)
+      assert retrieved_assignment.id == assignment.id
+    end
+
+    test "returns nil when equipment is not assigned", %{equipment: equipment} do
+      assert Agents.get_equipment_assignment(equipment.id) == nil
+    end
+  end
+
+  describe "update_equipment_assignment/2" do
+    setup %{organization: org} do
+      {:ok, site} =
+        Towerops.Sites.create_site(%{
+          name: "Test Site",
+          organization_id: org.id
+        })
+
+      {:ok, equipment} =
+        Towerops.Equipment.create_equipment(%{
+          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,
+        equipment: equipment,
+        agent_token1: agent_token1,
+        agent_token2: agent_token2
+      }
+    end
+
+    test "creates new assignment when equipment is unassigned", %{
+      equipment: equipment,
+      agent_token1: agent_token1
+    } do
+      assert {:ok, assignment} =
+               Agents.update_equipment_assignment(equipment.id, agent_token1.id)
+
+      assert assignment.agent_token_id == agent_token1.id
+      assert assignment.equipment_id == equipment.id
+    end
+
+    test "updates existing assignment to new agent", %{
+      equipment: equipment,
+      agent_token1: agent_token1,
+      agent_token2: agent_token2
+    } do
+      {:ok, _} = Agents.assign_equipment_to_agent(agent_token1.id, equipment.id)
+
+      assert {:ok, updated_assignment} =
+               Agents.update_equipment_assignment(equipment.id, agent_token2.id)
+
+      assert updated_assignment.agent_token_id == agent_token2.id
+      assert updated_assignment.equipment_id == equipment.id
+    end
+
+    test "removes assignment when set to nil", %{
+      equipment: equipment,
+      agent_token1: agent_token1
+    } do
+      {:ok, _} = Agents.assign_equipment_to_agent(agent_token1.id, equipment.id)
+
+      assert {:ok, nil} = Agents.update_equipment_assignment(equipment.id, nil)
+      assert Agents.get_equipment_assignment(equipment.id) == nil
+    end
+  end
+
   describe "list_agent_equipment/1" do
     setup %{organization: org} do
       {:ok, site} =