From 03fd2f1a4ee36debdb717a7c0d9329de081d0b0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marta=20G=C3=B3mez=20Mac=C3=ADas?= <mgmacias@google.com>
Date: Mon, 10 Jun 2024 12:41:26 +0200
Subject: [PATCH] fix(client): Fix Runtime error when running non async apps
 (#192)

---
 .github/workflows/testing.yaml |  2 +-
 setup.py                       |  4 +++-
 tests/test_client.py           | 28 +++++++++++++++++++++++++++-
 tests/wsgi_app.py              | 14 ++++++++++++++
 vt/client.py                   | 12 +++++++++++-
 5 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 tests/wsgi_app.py

diff --git a/.github/workflows/testing.yaml b/.github/workflows/testing.yaml
index 7b1d719..06918aa 100644
--- a/.github/workflows/testing.yaml
+++ b/.github/workflows/testing.yaml
@@ -19,7 +19,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pytest pylint aiohttp pytest_httpserver pytest_asyncio
+          pip install pytest pylint aiohttp pytest_httpserver pytest_asyncio flask
       - name: Lint
         run: |
           pylint --rcfile pylintrc vt/ tests examples/
diff --git a/setup.py b/setup.py
index 0b5fe96..75edb7b 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,9 @@
         "aiohttp ; python_version>'3.7'"
     ],
     setup_requires=["pytest-runner"],
-    extras_require={"test": ["pytest", "pytest_httpserver", "pytest_asyncio"]},
+    extras_require={
+        "test": ["pytest", "pytest_httpserver", "pytest_asyncio", "flask"]
+    },
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: Apache Software License",
diff --git a/tests/test_client.py b/tests/test_client.py
index a2ce5a8..6445463 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -15,6 +15,7 @@
 
 import datetime
 import io
+import functools
 import json
 import pickle
 
@@ -23,8 +24,10 @@
 from vt import Client
 from vt import Object
 
+from tests import wsgi_app
 
-def new_client(httpserver):
+
+def new_client(httpserver, unused_apikey=""):
   return Client(
       "dummy_api_key",
       host="http://" + httpserver.host + ":" + str(httpserver.port),
@@ -480,3 +483,26 @@ def test_user_headers(httpserver):
   assert "Accept-Encoding" in headers
   assert "User-Agent" in headers
   assert "foo" in headers
+
+
+def test_wsgi_app(httpserver, monkeypatch):
+  app = wsgi_app.app
+  app.config.update({"TESTING": True})
+  client = app.test_client()
+  expected_response = {"data": {
+      "id": "google.com",
+      "type": "domain",
+      "attributes": {"foo": "foo"},
+  }}
+
+
+  httpserver.expect_request(
+      "/api/v3/domains/google.com", method="GET",
+      headers={"X-Apikey": "dummy_api_key"}
+  ).respond_with_json(expected_response)
+  monkeypatch.setattr(
+      "tests.wsgi_app.vt.Client", functools.partial(new_client, httpserver)
+  )
+  response = client.get("/")
+  assert response.status_code == 200
+  assert response.json == expected_response
diff --git a/tests/wsgi_app.py b/tests/wsgi_app.py
new file mode 100644
index 0000000..35cc829
--- /dev/null
+++ b/tests/wsgi_app.py
@@ -0,0 +1,14 @@
+"""Simulates a Flask app using vt-py to test non-async environments."""
+
+import flask
+import os
+import vt
+
+app = flask.Flask(__name__)
+
+
+@app.get('/')
+def home():
+  with vt.Client(os.getenv('VT_APIKEY')) as c:
+    g = c.get('/domains/google.com')
+    return g.json()
diff --git a/vt/client.py b/vt/client.py
index 4fc7c5f..f5ed657 100644
--- a/vt/client.py
+++ b/vt/client.py
@@ -238,7 +238,17 @@ def __init__(
     if connector is not None:
       self._connector = connector
     else:
-      self._connector = aiohttp.TCPConnector(ssl=self._verify_ssl)
+      # the TCPConnector class expects to be instantiated inside a event loop.
+      # If there is none, create one.
+      try:
+        event_loop = asyncio.get_event_loop()
+      except RuntimeError:
+        event_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(event_loop)
+
+      self._connector = aiohttp.TCPConnector(
+          ssl=self._verify_ssl, loop=event_loop
+      )
 
   def _full_url(self, path:str, *args: typing.Any) -> str:
     try: