From ac6cd7be82f8bd8a127432f0d4a3ddadb971b062 Mon Sep 17 00:00:00 2001
From: Yury Shkoda <shkoda.yuri@gmail.com>
Date: Wed, 30 Jun 2021 16:55:09 +0300
Subject: [PATCH] fix(session): fixed polling session state

---
 src/SessionManager.ts        | 24 +++++++++---------------
 src/request/RequestClient.ts |  3 +++
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/src/SessionManager.ts b/src/SessionManager.ts
index 93a411d..914a686 100644
--- a/src/SessionManager.ts
+++ b/src/SessionManager.ts
@@ -6,10 +6,6 @@ import { RequestClient } from './request/RequestClient'
 const MAX_SESSION_COUNT = 1
 const RETRY_LIMIT: number = 3
 let RETRY_COUNT: number = 0
-const INTERNAL_SAS_ERROR = {
-  status: 304,
-  message: 'Not Modified'
-}
 
 export class SessionManager {
   constructor(
@@ -164,7 +160,7 @@ export class SessionManager {
 
     const stateLink = session.links.find((l: any) => l.rel === 'state')
 
-    return new Promise(async (resolve, _) => {
+    return new Promise(async (resolve, reject) => {
       if (
         sessionState === 'pending' ||
         sessionState === 'running' ||
@@ -182,7 +178,7 @@ export class SessionManager {
             etag!,
             accessToken
           ).catch((err) => {
-            throw err
+            throw prefixMessage(err, 'Error while getting session state.')
           })
 
           sessionState = state.trim()
@@ -196,13 +192,14 @@ export class SessionManager {
 
           // There is an internal error present in SAS Viya 3.5
           // Retry to wait for a session status in such case of SAS internal error
-          if (
-            sessionState === INTERNAL_SAS_ERROR.message &&
-            RETRY_COUNT < RETRY_LIMIT
-          ) {
-            RETRY_COUNT++
+          if (!sessionState) {
+            if (RETRY_COUNT < RETRY_LIMIT) {
+              RETRY_COUNT++
 
-            resolve(this.waitForSession(session, etag, accessToken))
+              resolve(this.waitForSession(session, etag, accessToken))
+            } else {
+              reject('Could not get session state.')
+            }
           }
 
           resolve(sessionState)
@@ -222,9 +219,6 @@ export class SessionManager {
       .get(url, accessToken, 'text/plain', { 'If-None-Match': etag })
       .then((res) => res.result as string)
       .catch((err) => {
-        if (err.status === INTERNAL_SAS_ERROR.status)
-          return INTERNAL_SAS_ERROR.message
-
         throw err
       })
   }
diff --git a/src/request/RequestClient.ts b/src/request/RequestClient.ts
index 580c1fa..c5e6e70 100644
--- a/src/request/RequestClient.ts
+++ b/src/request/RequestClient.ts
@@ -63,6 +63,9 @@ export class RequestClient implements HttpClient {
         baseURL: baseUrl
       })
     }
+
+    this.httpClient.defaults.validateStatus = (status) =>
+      status >= 200 && status < 305
   }
 
   public getCsrfToken(type: 'general' | 'file' = 'general') {