From f7d5778582f11ce804b6a03366467664eb3e200e Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Fri, 27 Sep 2024 11:13:47 -0400
Subject: [PATCH] set feature flags headers on collection httpClient as well

---
 src/driver/collection.ts | 4 +++-
 src/driver/connection.ts | 1 +
 tests/driver/api.test.ts | 5 ++++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/driver/collection.ts b/src/driver/collection.ts
index 1b34c75..2b42c20 100644
--- a/src/driver/collection.ts
+++ b/src/driver/collection.ts
@@ -50,7 +50,9 @@ export class Collection extends MongooseCollection {
 
     //getter for collection
     get collection() {
-        return this.conn.db.collection(this.name);
+        const collection = this.conn.db.collection(this.name);
+        Object.assign(collection._httpClient.baseHeaders, this.conn.featureFlags);
+        return collection;
     }
 
     /**
diff --git a/src/driver/connection.ts b/src/driver/connection.ts
index 8c3af14..70f966a 100644
--- a/src/driver/connection.ts
+++ b/src/driver/connection.ts
@@ -129,6 +129,7 @@ export class Connection extends MongooseConnection {
         const featureFlags: Record<string, 'true'>| null = Array.isArray(options && options.featureFlags)
             ? options.featureFlags.reduce((obj: Record<string, 'true'>, key: string) => Object.assign(obj, { [key]: 'true' }), {})
             : null;
+        this.featureFlags = featureFlags;
 
         const client = options?.isAstra
             ? new DataAPIClient(applicationToken)
diff --git a/tests/driver/api.test.ts b/tests/driver/api.test.ts
index 97a455a..4eaf51b 100644
--- a/tests/driver/api.test.ts
+++ b/tests/driver/api.test.ts
@@ -717,7 +717,6 @@ describe('Mongoose Model API level tests', async () => {
             assert.ok(res.status.collections.includes('carts'));
         });
         it('API ops tests collection.runCommand()', async () => {
-            console.log('XT', mongooseInstance!.connection.db.collection('carts'));
             const res = await mongooseInstance!.connection.db.collection('carts')._httpClient.executeCommand({ find: {} });
             assert.ok(Array.isArray(res.data.documents));
         });
@@ -752,6 +751,10 @@ describe('Mongoose Model API level tests', async () => {
                 }
             });
             assert.ok(res.status.ok);
+
+            const Bot = mongoose.model('Bot', new mongoose.Schema({ name: String }), 'bots');
+            await Bot.findOne();
+
             await mongoose.disconnect();
         });
     });