Pytest initial setup (#79)
* add pytest and logic to test connectors * Readme to run pytest * precommit fixes
taylor_socfortress committed
Jul 27, 2023 at 09:25 UTC
094c1edfe0a3e0e57ebb2aa0aa6dc6b98c6292d0
5 files changed
+149
README.md
+14
@@ -93,6 +93,20 @@ Don't forget to commit the new migration file into git.
93
94
See https://flask-migrate.readthedocs.io/en/latest/ for further information
95
96
+## Running Pytest
97
+
98
+To run pytest, ensure pytest is installed:
99
+
100
+```
101
+pip install pytest
102
+```
103
+
104
+Then run pytest:
105
+
106
+```
107
+pytest .\backend\tests\routes\test_connectors.py
108
+```
109
+
110
# Deployment
111
112
## Production Deployment Notes
backend/requirements.in
+1
@@ -16,6 +16,7 @@ mitreattack-python
16
openai
17
pika
18
psycopg2-binary
19
+pytest
20
pyvelociraptor~=0.1
21
reportlab
22
requests
backend/tests/__init__.py
backend/tests/routes/__init__.py
backend/tests/routes/test_connectors.py
new
+134
@@ -0,0 +1,134 @@
1
+import json
2
+
3
+import pytest
4
+
5
+from app import db # Import the SQLAlchemy instance from your application
6
+
7
+# Add your connector classes here
8
+from app.models.connectors import AskSOCFortressConnector
9
+from app.models.connectors import ConnectorFactory
10
+from app.models.connectors import DfirIrisConnector
11
+from app.models.connectors import GraylogConnector
12
+from app.models.connectors import InfluxDBConnector
13
+from app.models.connectors import RabbitMQConnector
14
+from app.models.connectors import ShuffleConnector
15
+from app.models.connectors import SocfortressThreatIntelConnector
16
+from app.models.connectors import SublimeConnector
17
+from app.models.connectors import VelociraptorConnector
18
+from app.models.connectors import WazuhIndexerConnector
19
+from app.models.connectors import WazuhManagerConnector
20
+from app.models.models import Connectors
21
+from copilot import app # Import your Flask application creation function
22
+
23
+CONNECTORS = {
24
+ "Wazuh-Indexer": WazuhIndexerConnector,
25
+ "Graylog": GraylogConnector,
26
+ "Wazuh-Manager": WazuhManagerConnector,
27
+ "DFIR-IRIS": DfirIrisConnector,
28
+ "Velociraptor": VelociraptorConnector,
29
+ "RabbitMQ": RabbitMQConnector,
30
+ "Shuffle": ShuffleConnector,
31
+ "Sublime": SublimeConnector,
32
+ "InfluxDB": InfluxDBConnector,
33
+ "AskSocfortress": AskSOCFortressConnector,
34
+ "SocfortressThreatIntel": SocfortressThreatIntelConnector,
35
+}
36
+
37
+
38
+@pytest.fixture(scope="function")
39
+def setup_connector_factory():
40
+ connector_factory = ConnectorFactory()
41
+ for name, cls in CONNECTORS.items():
42
+ connector_factory.register_creator(name, cls)
43
+ return connector_factory
44
+
45
+
46
+@pytest.fixture(scope="function")
47
+def client():
48
+ app.config["TESTING"] = True
49
+ with app.test_client() as testing_client:
50
+ app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///test_copilot.db" # Use a separate SQLite database for testing
51
+ with app.app_context():
52
+ # Create the database tables
53
+ db.create_all()
54
+
55
+ # Insert your test data here
56
+ test_connector = Connectors(
57
+ connector_name="Wazuh-Indexer",
58
+ connector_type="Test Type",
59
+ connector_url="http://localhost:8000",
60
+ connector_username="testuser",
61
+ connector_password="testpass",
62
+ connector_api_key="testapikey",
63
+ )
64
+
65
+ db.session.add(test_connector)
66
+ db.session.commit()
67
+
68
+ yield testing_client
69
+
70
+ # Tear down the database after the tests
71
+ db.drop_all()
72
+
73
+
74
+def test_list_connectors_available(client):
75
+ response = client.get("/connectors")
76
+ assert response.status_code == 200
77
+
78
+ # The response data is a byte string so we need to decode it
79
+ data = json.loads(response.data.decode())
80
+
81
+ # Check that the response contains the expected keys
82
+ assert "message" in data
83
+ assert "connectors" in data
84
+ assert "success" in data
85
+
86
+ # Check that the 'success' field is True
87
+ assert data["success"] is True
88
+
89
+
90
+def test_get_connector_details(client, setup_connector_factory):
91
+ # Use setup_connector_factory here
92
+ # Define the id of the connector you want to test
93
+ id = 1 # Modify this value based on your test data
94
+
95
+ # Send a GET request to the endpoint
96
+ response = client.get(f"/connectors/{id}")
97
+
98
+ # Convert the response data from json to a Python dictionary
99
+ data = json.loads(response.data)
100
+
101
+ # Assert that the request got a success response (HTTP status code: 200)
102
+ assert response.status_code == 200
103
+
104
+ # Assert that the data has these keys
105
+ assert "message" in data
106
+ assert "connector" in data
107
+ assert "success" in data
108
+
109
+ # Assert that success is True
110
+ assert data["success"] is True
111
+
112
+
113
+def test_get_connector_details_error(client):
114
+ # Define the id of the connector you want to test
115
+ id = 999999 # This should be an ID that does not exist in your test database
116
+
117
+ # Send a GET request to the endpoint
118
+ response = client.get(f"/connectors/{id}")
119
+
120
+ # Convert the response data from json to a Python dictionary
121
+ data = json.loads(response.data)
122
+
123
+ # Assert that the request got a not found response (HTTP status code: 404)
124
+ assert response.status_code == 404
125
+
126
+ # Assert that the data has these keys
127
+ assert "message" in data
128
+ assert "success" in data
129
+
130
+ # Assert that success is False
131
+ assert data["success"] is False
132
+
133
+ # Assert that the message is "Connector not found"
134
+ assert data["message"] == "Connector not found"