From 98556fac9dd12024d52bb19ba0c09b098c6e8d64 Mon Sep 17 00:00:00 2001
From: Jhonathan Abreu <jdabreu25@gmail.com>
Date: Tue, 20 Feb 2024 16:14:32 -0400
Subject: [PATCH] Minor fix in credentials being overriden by SetJob (#7)

---
 .../PolygonDataQueueHandlerCommonTests.cs     | 20 +++++++++++++++++++
 .../PolygonDataQueueHandler.cs                | 14 ++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/QuantConnect.Polygon.Tests/PolygonDataQueueHandlerCommonTests.cs b/QuantConnect.Polygon.Tests/PolygonDataQueueHandlerCommonTests.cs
index 1ef1f1e..44be5f6 100644
--- a/QuantConnect.Polygon.Tests/PolygonDataQueueHandlerCommonTests.cs
+++ b/QuantConnect.Polygon.Tests/PolygonDataQueueHandlerCommonTests.cs
@@ -111,6 +111,26 @@ public void CanInitializeUsingJobPacket()
             Config.Set("polygon-api-key", _apiKey);
         }
 
+        [Test]
+        public void JobPacketWontOverrideCredentials()
+        {
+            var job = new LiveNodePacket
+            {
+                BrokerageData = new Dictionary<string, string>() { { "polygon-api-key", "InvalidApiKeyThatWontBeUsed" } }
+            };
+
+            // This should pick up the API key from the configuration
+            using var polygon = new PolygonDataQueueHandler();
+
+            Assert.IsFalse(polygon.IsConnected);
+
+            // This should not override the API key
+            polygon.SetJob(job);
+
+            // Since API key was not overriden, this should work
+            AssertConnection(polygon);
+        }
+
         private SubscriptionDataConfig GetSubscriptionDataConfig<T>(Symbol symbol, Resolution resolution)
         {
             return new SubscriptionDataConfig(
diff --git a/QuantConnect.Polygon/PolygonDataQueueHandler.cs b/QuantConnect.Polygon/PolygonDataQueueHandler.cs
index 049c55b..b829809 100644
--- a/QuantConnect.Polygon/PolygonDataQueueHandler.cs
+++ b/QuantConnect.Polygon/PolygonDataQueueHandler.cs
@@ -63,6 +63,7 @@ public partial class PolygonDataQueueHandler : IDataQueueHandler
         private readonly MarketHoursDatabase _marketHoursDatabase = MarketHoursDatabase.FromDataFolder();
         private readonly Dictionary<Symbol, DateTimeZone> _symbolExchangeTimeZones = new();
 
+        private bool _initialized;
         private bool _disposed;
 
         private bool _unsupportedSecurityTypeMessageLogged;
@@ -130,6 +131,7 @@ public PolygonDataQueueHandler(string apiKey, int maxSubscriptionsPerWebSocket,
         /// </summary>
         private void Initialize(int maxSubscriptionsPerWebSocket, bool streamingEnabled = true)
         {
+            _initialized = true;
             _dataAggregator = new PolygonAggregationManager();
             RestApiClient = new PolygonRestApiClient(_apiKey);
             _optionChainProvider = new CachingOptionChainProvider(new PolygonOptionChainProvider(RestApiClient, _symbolMapper));
@@ -163,12 +165,18 @@ private void Initialize(int maxSubscriptionsPerWebSocket, bool streamingEnabled
         /// <param name="job">Job we're subscribing for</param>
         public void SetJob(LiveNodePacket job)
         {
-            job.BrokerageData.TryGetValue("polygon-api-key", out _apiKey);
-            if (string.IsNullOrEmpty(_apiKey))
+            if (_initialized)
             {
-                throw new ArgumentException("The PolygonDataQueueHandler requires an API key");
+                return;
+            }
+
+            if (!job.BrokerageData.TryGetValue("polygon-api-key", out var apiKey))
+            {
+                throw new ArgumentException("The Polygon.io API key is missing from the brokerage data.");
             }
 
+            _apiKey = apiKey;
+
             var maxSubscriptionsPerWebSocket = 0;
             if (!job.BrokerageData.TryGetValue("polygon-max-subscriptions-per-websocket", out var maxSubscriptionsPerWebSocketStr) ||
                 !int.TryParse(maxSubscriptionsPerWebSocketStr, out maxSubscriptionsPerWebSocket))