reduce time for expiration test

Author: udaij12

docs/token_authorization_api.md

@@ -1,7 +1,7 @@
 # TorchServe token authorization API
 
 ## Configuration
-1. Enable token authorization by adding the provided plugin at start using the `--plugin-path` command.
+1. Enable token authorization by adding the provided plugin at start using the `--plugins-path` command.
 2. Torchserve will enable token authorization if the plugin is provided. In the current working directory a file `key_file.json` will be generated.
     1. Example key file:
 

frontend/server/src/main/java/org/pytorch/serve/http/TokenAuthorizationHandler.java

@@ -26,7 +26,7 @@ public class TokenAuthorizationHandler extends HttpRequestHandlerChain {
     private static Boolean tokenEnabled = false;
     private static Class<?> tokenClass;
     private static Object tokenObject;
-    private static Integer timeToExpirationMinutes = 60;
+    private static Double timeToExpirationMinutes = 60.0;
 
     /** Creates a new {@code InferenceRequestHandler} instance. */
     public TokenAuthorizationHandler(TokenType type) {
@@ -59,9 +59,9 @@ public static void setupTokenClass() {
         try {
             tokenClass = Class.forName("org.pytorch.serve.plugins.endpoint.Token");
             tokenObject = tokenClass.getDeclaredConstructor().newInstance();
-            Method method = tokenClass.getMethod("setTime", Integer.class);
-            Integer time = ConfigManager.getInstance().getTimeToExpiration();
-            if (time != 0) {
+            Method method = tokenClass.getMethod("setTime", Double.class);
+            Double time = ConfigManager.getInstance().getTimeToExpiration();
+            if (time != 0.0) {
                 timeToExpirationMinutes = time;
             }
             method.invoke(tokenObject, timeToExpirationMinutes);

frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java

@@ -860,15 +860,15 @@ public boolean isSnapshotDisabled() {
         return snapshotDisabled;
     }
 
-    public Integer getTimeToExpiration() {
+    public Double getTimeToExpiration() {
         if (prop.getProperty(TS_TOKEN_EXPIRATION_TIME_MIN) != null) {
             try {
-                return Integer.valueOf(prop.getProperty(TS_TOKEN_EXPIRATION_TIME_MIN));
+                return Double.valueOf(prop.getProperty(TS_TOKEN_EXPIRATION_TIME_MIN));
             } catch (NumberFormatException e) {
                 logger.error("Token expiration not a valid integer");
             }
         }
-        return 0;
+        return 0.0;
     }
 
     public boolean isSSLEnabled(ConnectorType connectorType) {

plugins/endpoints/src/main/java/org/pytorch/serve/plugins/endpoint/Token.java

@@ -19,7 +19,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.TimeUnit;
 import org.pytorch.serve.servingsdk.Context;
 import org.pytorch.serve.servingsdk.ModelServerEndpoint;
 import org.pytorch.serve.servingsdk.annotations.Endpoint;
@@ -39,7 +38,7 @@ public class Token extends ModelServerEndpoint {
     private static String inferenceKey;
     private static Instant managementExpirationTimeMinutes;
     private static Instant inferenceExpirationTimeMinutes;
-    private static Integer timeToExpirationMinutes;
+    private static Double timeToExpirationMinutes;
     private SecureRandom secureRandom = new SecureRandom();
     private Base64.Encoder baseEncoder = Base64.getUrlEncoder();
     private String fileName = "key_file.json";
@@ -80,7 +79,8 @@ public String generateKey() {
     }
 
     public Instant generateTokenExpiration() {
-        return Instant.now().plusSeconds(TimeUnit.MINUTES.toSeconds(timeToExpirationMinutes));
+        long secondsToAdd = (long) (timeToExpirationMinutes * 60);
+        return Instant.now().plusSeconds(secondsToAdd);
     }
 
     // generates a key file with new keys depending on the parameter provided
@@ -217,7 +217,7 @@ public Instant getManagementExpirationTime() {
         return managementExpirationTimeMinutes;
     }
 
-    public void setTime(Integer time) {
+    public void setTime(Double time) {
         timeToExpirationMinutes = time;
     }
 }

test/pytest/test_token_authorization.py

@@ -66,6 +66,8 @@ def setup_torchserve():
     MODEL_STORE = os.path.join(ROOT_DIR, "model_store/")
     PLUGIN_STORE = os.path.join(ROOT_DIR, "plugins-path")
 
+    Path(test_utils.MODEL_STORE).mkdir(parents=True, exist_ok=True)
+
     test_utils.start_torchserve(no_config_snapshots=True, plugin_folder=PLUGIN_STORE)
 
     key = read_key_file("management")
@@ -94,10 +96,10 @@ def setup_torchserve_expiration():
     MODEL_STORE = os.path.join(ROOT_DIR, "model_store/")
     PLUGIN_STORE = os.path.join(ROOT_DIR, "plugins-path")
 
+    Path(test_utils.MODEL_STORE).mkdir(parents=True, exist_ok=True)
+
     test_utils.start_torchserve(
-        snapshot_file=config_file,
-        no_config_snapshots=True,
-        plugin_folder=PLUGIN_STORE,
+        snapshot_file=config_file, no_config_snapshots=True, plugin_folder=PLUGIN_STORE
     )
 
     key = read_key_file("management")
@@ -214,13 +216,14 @@ def test_token_management_api(setup_torchserve):
 
 
 # Test expiration time
+@pytest.mark.module2
 def test_token_expiration_time(setup_torchserve_expiration):
     key = read_key_file("management")
     header = {"Authorization": f"Bearer {key}"}
     response = requests.get("http://localhost:8081/models/resnet18", headers=header)
     assert response.status_code == 200, "Token check failed"
 
-    time.sleep(60)
+    time.sleep(15)
 
     response = requests.get("http://localhost:8081/models/resnet18", headers=header)
     assert response.status_code == 400, "Token check failed"

test/resources/config_token.properties

@@ -1 +1 @@
-token_expiration_min=1
+token_expiration_min=0.25

ts/model_server.py

@@ -3,6 +3,7 @@
 """
 
 import os
+import pathlib
 import platform
 import re
 import subprocess
@@ -48,7 +49,7 @@ def start() -> None:
             try:
                 parent = psutil.Process(pid)
                 parent.terminate()
-                os.remove(os.getcwd() + "/key_file.json")
+                pathlib.Path("key_file.json").unlink(missing_ok=True)
                 if args.foreground:
                     try:
                         parent.wait(timeout=60)