diff --git a/channel_test.go b/channel_test.go
index bb60653..87eb2d6 100644
--- a/channel_test.go
+++ b/channel_test.go
@@ -411,7 +411,7 @@ func TestChannel_SendRestrictedVisibilityMessage(t *testing.T) {
 
 	resp, err := ch.SendMessage(ctx, msg, adminUser.ID)
 	require.NoError(t, err, "send message")
-	assert.Equal(t, resp.Message.RestrictedVisibility, msg.RestrictedVisibility)
+	assert.Equal(t, msg.RestrictedVisibility, resp.Message.RestrictedVisibility)
 }
 
 func TestChannel_Truncate(t *testing.T) {
diff --git a/message_test.go b/message_test.go
index bc61551..89404b5 100644
--- a/message_test.go
+++ b/message_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -146,3 +147,56 @@ func TestClient_SendMessage_KeepChannelHidden(t *testing.T) {
 	require.NoError(t, err)
 	require.Empty(t, result.Channels)
 }
+
+func TestClient_UpdateRestrictedVisibilityMessage(t *testing.T) {
+	c := initClient(t)
+	ch := initChannel(t, c)
+	ctx := context.Background()
+	adminUser := randomUserWithRole(t, c, "admin")
+	user1 := randomUser(t, c)
+	user2 := randomUser(t, c)
+	msg := &Message{
+		Text: "test message",
+		RestrictedVisibility: []string{
+			user1.ID,
+		},
+	}
+
+	resp, err := ch.SendMessage(ctx, msg, adminUser.ID)
+	require.NoError(t, err, "send message")
+
+	msg = resp.Message
+	msg.RestrictedVisibility = []string{user2.ID}
+	resp, err = c.UpdateMessage(ctx, msg, msg.ID)
+	require.NoError(t, err, "send message")
+	assert.Equal(t, []string{user2.ID}, resp.Message.RestrictedVisibility)
+}
+
+func TestClient_PartialUpdateRestrictedVisibilityMessage(t *testing.T) {
+	c := initClient(t)
+	ch := initChannel(t, c)
+	ctx := context.Background()
+	adminUser := randomUserWithRole(t, c, "admin")
+	user1 := randomUser(t, c)
+	user2 := randomUser(t, c)
+	msg := &Message{
+		Text: "test message",
+		RestrictedVisibility: []string{
+			user1.ID,
+		},
+	}
+
+	_, err := ch.SendMessage(ctx, msg, adminUser.ID)
+	require.NoError(t, err, "send message")
+
+	resp, err := c.PartialUpdateMessage(ctx, msg.ID, &MessagePartialUpdateRequest{
+		UserID: adminUser.ID,
+		PartialUpdate: PartialUpdate{
+			Set: map[string]interface{}{
+				"restricted_visibility": []string{user2.ID},
+			},
+		},
+	})
+	require.NoError(t, err, "send message")
+	assert.Equal(t, []string{user2.ID}, resp.Message.RestrictedVisibility)
+}