diff --git a/tests/integration/IntegrationTest.php b/tests/integration/IntegrationTest.php
index bd5ce8b..0c7e8be 100644
--- a/tests/integration/IntegrationTest.php
+++ b/tests/integration/IntegrationTest.php
@@ -217,6 +217,21 @@ public function testUpsertUser()
         $this->client->deleteUser($user["id"], ["user" => "hard", "messages" => "hard"]);
     }
 
+    public function testUpsertUserWithTeam()
+    {
+        $user = [
+            "id" => $this->generateGuid(),
+            "team" => "blue",
+            "teams_role" => ["blue" => "admin"]
+        ];
+        $response = $this->client->upsertUser($user);
+        $this->assertTrue(array_key_exists("users", (array)$response));
+        $this->assertTrue(array_key_exists($user["id"], $response["users"]));
+        $this->assertEquals("blue", $response["users"][$user["id"]]["team"]);
+        $this->assertEquals("admin", $response["users"][$user["id"]]["teams_role"]["blue"]);
+        $this->client->deleteUser($user["id"], ["user" => "hard", "messages" => "hard"]);
+    }
+
     public function testUpsertUsers()
     {
         $user = ["id" => $this->generateGuid()];
@@ -1188,6 +1203,31 @@ public function testPartialUpdateUsers()
         $this->client->deleteUser($wally["id"], ["user" => "hard", "messages" => "hard"]);
     }
 
+    public function testPartialUpdateUserWithTeam()
+    {
+        $user = ["id" => $this->generateGuid(), "name" => "Test User"];
+        $response = $this->client->upsertUser($user);
+        $this->assertTrue(array_key_exists("users", (array)$response));
+        $this->assertTrue(array_key_exists($user["id"], $response["users"]));
+
+        // Partially update the user with team and teams_role
+        $response = $this->client->partialUpdateUser([
+            "id" => $user["id"],
+            "set" => [
+                "team" => "blue",
+                "teams_role" => ["blue" => "admin"]
+            ]
+        ]);
+
+        // Verify the changes
+        $response = $this->client->queryUsers(["id" => $user["id"]]);
+        $this->assertEquals("blue", $response["users"][0]["team"]);
+        $this->assertEquals("admin", $response["users"][0]["teams_role"]["blue"]);
+
+        // Clean up
+        $this->client->deleteUser($user["id"], ["user" => "hard", "messages" => "hard"]);
+    }
+
     public function testChannelMuteUnmute()
     {
         // setup