diff --git a/apps/chat-e2e/src/assertions/manageAttachmentsAssertion.ts b/apps/chat-e2e/src/assertions/manageAttachmentsAssertion.ts
index dfb55b06b..aaf837d33 100644
--- a/apps/chat-e2e/src/assertions/manageAttachmentsAssertion.ts
+++ b/apps/chat-e2e/src/assertions/manageAttachmentsAssertion.ts
@@ -79,10 +79,15 @@ export class ManageAttachmentsAssertion {
     state: ElementCaretState,
   ) {
     const sectionElement = this.attachFilesModal.getSectionElement(section);
-    const isExpanded =
-      await this.attachFilesModal.isSectionExpanded(sectionElement);
+    const filesSection = this.attachFilesModal.getFilesSection(sectionElement);
     state === 'expanded'
-      ? expect(isExpanded, `Section "${section}" is ${state}`).toBeTruthy()
-      : expect(isExpanded, `Section "${section}" is ${state}`).toBeFalsy();
+      ? await expect(
+          filesSection,
+          `Section "${section}" is ${state}`,
+        ).toBeVisible()
+      : await expect(
+          filesSection,
+          `Section "${section}" is ${state}`,
+        ).toBeHidden();
   }
 }
diff --git a/apps/chat-e2e/src/ui/webElements/attachFilesModal.ts b/apps/chat-e2e/src/ui/webElements/attachFilesModal.ts
index 3ebe32ce8..d0472b233 100644
--- a/apps/chat-e2e/src/ui/webElements/attachFilesModal.ts
+++ b/apps/chat-e2e/src/ui/webElements/attachFilesModal.ts
@@ -163,6 +163,11 @@ export class AttachFilesModal extends BaseElement {
     SelectFolderModalSelectors.newFolderButton,
   );
 
+  public getFilesSection = (sectionElement: BaseElement) =>
+    sectionElement
+      .getChildElementBySelector(AttachFilesModalSelectors.fileSection)
+      .getElementLocator();
+
   public closeButton = this.getChildElementBySelector(IconSelectors.cancelIcon);
 
   public async checkAttachedFile(
@@ -209,12 +214,4 @@ export class AttachFilesModal extends BaseElement {
       ErrorLabelSelectors.errorText,
     ).getElementContent();
   }
-
-  public async isSectionExpanded(
-    sectionElement: BaseElement,
-  ): Promise<boolean> {
-    return sectionElement
-      .getChildElementBySelector(AttachFilesModalSelectors.fileSection)
-      .isVisible();
-  }
 }