From 2939cde23862f34df67016e3d4f01386111f84b4 Mon Sep 17 00:00:00 2001
From: Serge Klochkov <3175289+slvrtrn@users.noreply.github.com>
Date: Tue, 22 Oct 2024 13:57:34 +0200
Subject: [PATCH] Add read-only users cleanup to the tests (#343)

---
 .../integration/read_only_user.test.ts        | 19 ++++++---
 .../integration/web_abort_request.test.ts     | 39 ++++++++++++-------
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/packages/client-common/__tests__/integration/read_only_user.test.ts b/packages/client-common/__tests__/integration/read_only_user.test.ts
index 260fed26..73fd9b06 100644
--- a/packages/client-common/__tests__/integration/read_only_user.test.ts
+++ b/packages/client-common/__tests__/integration/read_only_user.test.ts
@@ -1,17 +1,21 @@
 import type { ClickHouseClient } from '@clickhouse/client-common'
+import { isCloudTestEnv } from '@test/utils/test_env'
 import { createReadOnlyUser } from '../fixtures/read_only_user'
 import { createSimpleTable } from '../fixtures/simple_table'
 import { createTestClient, getTestDatabaseName, guid } from '../utils'
 
 describe('read only user', () => {
+  let defaultClient: ClickHouseClient
   let client: ClickHouseClient
   let tableName: string
+  let userName: string
 
   beforeAll(async () => {
     const database = getTestDatabaseName()
-    const defaultClient = createTestClient()
+    defaultClient = createTestClient()
 
-    const { username, password } = await createReadOnlyUser(defaultClient)
+    const credentials = await createReadOnlyUser(defaultClient)
+    userName = credentials.username
 
     // Populate some test table to select from
     tableName = `read_only_user_data_${guid()}`
@@ -20,13 +24,12 @@ describe('read only user', () => {
       table: tableName,
       values: [[42, 'hello', [0, 1]]],
     })
-    await defaultClient.close()
 
     // Create a client that connects read only user to the test database
     client = createTestClient({
-      username,
-      password,
       database,
+      username: credentials.username,
+      password: credentials.password,
       clickhouse_settings: {
         // readonly user cannot adjust settings. reset the default ones set by fixtures.
         // might be fixed by https://github.com/ClickHouse/ClickHouse/issues/40244
@@ -37,7 +40,13 @@ describe('read only user', () => {
   })
 
   afterAll(async () => {
+    if (isCloudTestEnv()) {
+      await defaultClient.command({
+        query: `DROP USER IF EXISTS ${userName}`,
+      })
+    }
     await client.close()
+    await defaultClient.close()
   })
 
   it('should select some data without issues', async () => {
diff --git a/packages/client-web/__tests__/integration/web_abort_request.test.ts b/packages/client-web/__tests__/integration/web_abort_request.test.ts
index 426fee70..dee28bbf 100644
--- a/packages/client-web/__tests__/integration/web_abort_request.test.ts
+++ b/packages/client-web/__tests__/integration/web_abort_request.test.ts
@@ -1,11 +1,12 @@
-import type { ClickHouseClient, Row } from '@clickhouse/client-common'
+import type { Row } from '@clickhouse/client-common'
 import { createTestClient } from '@test/utils'
+import type { WebClickHouseClient } from '../../src/client'
 
-describe('[Web] abort request', () => {
-  let client: ClickHouseClient<ReadableStream>
+fdescribe('[Web] abort request', () => {
+  let client: WebClickHouseClient
 
   beforeEach(() => {
-    client = createTestClient()
+    client = createTestClient() as unknown as WebClickHouseClient
   })
 
   afterEach(async () => {
@@ -31,23 +32,27 @@ describe('[Web] abort request', () => {
 
   it('cancels a select query while reading response', async () => {
     const controller = new AbortController()
+    let rowCount = 0
     const selectPromise = client
       .query({
-        query: 'SELECT * from system.numbers LIMIT 100000',
+        query: 'SELECT number FROM system.numbers LIMIT 1000',
         format: 'JSONCompactEachRow',
         abort_signal: controller.signal,
+        clickhouse_settings: {
+          // low block size to force streaming 1 row at a time
+          max_block_size: '1',
+        },
       })
       .then(async (rs) => {
         const reader = rs.stream().getReader()
         while (true) {
           const { done, value: rows } = await reader.read()
           if (done) break
-          ;(rows as Row[]).forEach((row: Row) => {
-            const [number] = row.json<[string]>()
-            // abort when reach number 3
-            if (number === '3') {
+          ;(rows as Row[]).forEach(() => {
+            if (rowCount >= 1) {
               controller.abort()
             }
+            rowCount++
           })
         }
       })
@@ -61,10 +66,15 @@ describe('[Web] abort request', () => {
   })
 
   it('cancels a select query while reading response by closing response stream', async () => {
+    let rowCount = 0
     const selectPromise = client
       .query({
-        query: 'SELECT * from system.numbers LIMIT 100000',
+        query: 'SELECT number FROM system.numbers LIMIT 3',
         format: 'JSONCompactEachRow',
+        clickhouse_settings: {
+          // low block size to force streaming 1 row at a time
+          max_block_size: '1',
+        },
       })
       .then(async function (rs) {
         const reader = rs.stream().getReader()
@@ -72,12 +82,11 @@ describe('[Web] abort request', () => {
           const { done, value: rows } = await reader.read()
           if (done) break
           for (const row of rows as Row[]) {
-            const [number] = row.json<[string]>()
-            // abort when reach number 3
-            if (number === '3') {
-              await reader.releaseLock()
-              await rs.close()
+            row.json()
+            if (rowCount >= 1) {
+              await rs.stream().cancel()
             }
+            rowCount++
           }
         }
       })