add test

Author: Sita04

genai/text_generation/test_text_generation.py

@@ -31,13 +31,17 @@
 import textgen_with_txt_stream
 import textgen_with_video
 import textgen_with_youtube_video
+import txt_gen 
 
 
 os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
 os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
 # The project name is included in the CICD pipeline
 # os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
 
+def test_txt_gen() -> None:
+    response = txt_gen.generate_content()
+    assert response
 
 def test_textgen_with_txt_stream() -> None:
     response = textgen_with_txt_stream.generate_content()

genai/text_generation/txt_gen.py

@@ -1,13 +1,15 @@
-from google import genai
 
-client = genai.Client(http_options={'api_version': 'v1'})
-response = client.models.generate_content(
-    model="gemini-2.0-flash-001",
-    contents="How does AI work?"
-)
-print(response.text)
-# Example response:
-# Okay, let's break down how AI works. It's a broad field, so I'll focus on the ...
-#
-# Here's a simplified overview:
-# ...
+def generate_content():
+    from google import genai
+
+    client = genai.Client(http_options={'api_version': 'v1'})
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents="How does AI work?"
+    )
+    return response.text
+    # Example response:
+    # Okay, let's break down how AI works. It's a broad field, so I'll focus on the ...
+    #
+    # Here's a simplified overview:
+    # ...