From ad42fcd0dabed13072f41d2abf04dabcd8505a57 Mon Sep 17 00:00:00 2001
From: Shyam Sai <shyam.s.sai@gmail.com>
Date: Fri, 20 Dec 2024 16:43:53 -0600
Subject: [PATCH] Fix style

---
 .../services/openai/test_OpenAIDefaults.py    | 23 +++++++++++--------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/cognitive/src/test/python/synapsemltest/services/openai/test_OpenAIDefaults.py b/cognitive/src/test/python/synapsemltest/services/openai/test_OpenAIDefaults.py
index a24190c3dc..beb86c49bb 100644
--- a/cognitive/src/test/python/synapsemltest/services/openai/test_OpenAIDefaults.py
+++ b/cognitive/src/test/python/synapsemltest/services/openai/test_OpenAIDefaults.py
@@ -3,7 +3,7 @@
 
 from synapse.ml.services.openai.OpenAIDefaults import OpenAIDefaults
 from synapse.ml.services.openai.OpenAIPrompt import OpenAIPrompt
-import unittest,os, json, subprocess
+import unittest, os, json, subprocess
 from pyspark.sql import SQLContext
 from pyspark.sql.functions import col
 
@@ -74,11 +74,14 @@ def test_prompt_w_defaults(self):
         )
         openai_api_key = json.loads(secretJson)["value"]
 
-        df = spark.createDataFrame([
-            ("apple", "fruits"),
-            ("mercedes", "cars"),
-            ("cake", "dishes"),
-        ], ["text", "category"])
+        df = spark.createDataFrame(
+            [
+                ("apple", "fruits"),
+                ("mercedes", "cars"),
+                ("cake", "dishes"),
+            ],
+            ["text", "category"],
+        )
 
         defaults = OpenAIDefaults()
         defaults.set_deployment_name("gpt-35-turbo-0125")
@@ -88,11 +91,13 @@ def test_prompt_w_defaults(self):
 
         prompt = OpenAIPrompt()
         prompt = prompt.setOutputCol("outParsed")
-        prompt = prompt.setPromptTemplate("Complete this comma separated list of 5 {category}: {text}, ")
+        prompt = prompt.setPromptTemplate(
+            "Complete this comma separated list of 5 {category}: {text}, "
+        )
         results = prompt.transform(df)
-        results.select("outParsed").show(truncate = False)
+        results.select("outParsed").show(truncate=False)
         nonNullCount = results.filter(col("outParsed").isNotNull()).count()
-        assert (nonNullCount == 3)
+        assert nonNullCount == 3
 
 
 if __name__ == "__main__":