Skip to content

Commit 0b2727a

Browse files
tomchoptomchop
committed
Add tests
1 parent fe1448f commit 0b2727a

File tree

6 files changed

+323
-9
lines changed

6 files changed

+323
-9
lines changed

‎.github/workflows/unittests.yml‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
name: Unit tests
2+
3+
on: [pull_request]
4+
5+
jobs:
6+
unittest:
7+
runs-on: ubuntu-latest
8+
strategy:
9+
matrix:
10+
os: [ubuntu-latest]
11+
python-version: ["3.10"]
12+
steps:
13+
- uses: actions/checkout@v4
14+
- run:
15+
sudo apt-get update && sudo apt-get install -y python3-pip && sudo
16+
pip3 install poetry
17+
- name: Set up Python ${{ matrix.python-version }}
18+
uses: actions/setup-python@v4
19+
with:
20+
python-version: ${{ matrix.python-version }}
21+
cache: poetry
22+
- name: Install Python dependencies
23+
run: poetry install --no-root
24+
- name: Test with unittest
25+
run: |
26+
poetry run python -m unittest discover -s tests/ -p '*.py'

‎.vscode/settings.json‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
{
2+
"python.testing.unittestArgs": [
3+
"-v",
4+
"-s",
5+
"./tests",
6+
"-p",
7+
"*.py"
8+
],
9+
"python.testing.pytestEnabled": false,
10+
"python.testing.unittestEnabled": true
11+
}

‎poetry.lock‎

Lines changed: 16 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ packages = [
1313
python = "^3.10"
1414
requests = "^2.32.3"
1515
click = "^8.1.7"
16+
requests-toolbelt = "^1.0.0"
1617

1718

1819
[tool.poetry.group.dev.dependencies]

‎tests/api.py‎

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
import unittest
2+
from unittest.mock import patch, MagicMock
3+
from yeti.api import YetiApi
4+
5+
6+
class TestYetiApi(unittest.TestCase):
7+
def setUp(self):
8+
self.api = YetiApi("http://fake-url")
9+
10+
@patch("yeti.api.requests.Session.post")
11+
def test_auth_api_key(self, mock_post):
12+
mock_response = MagicMock()
13+
mock_response.text = '{"access_token": "fake_token"}'
14+
mock_post.return_value = mock_response
15+
16+
self.api.auth_api_key("fake_apikey")
17+
self.assertEqual(self.api.client.headers["authorization"], "Bearer fake_token")
18+
mock_post.assert_called_with(
19+
"http://fake-url/api/v2/auth/api-token",
20+
headers={"x-yeti-apikey": "fake_apikey"},
21+
)
22+
23+
@patch("yeti.api.requests.Session.post")
24+
def test_search_indicators(self, mock_post):
25+
mock_response = MagicMock()
26+
mock_response.json.return_value = {"indicators": [{"name": "test"}]}
27+
mock_post.return_value = mock_response
28+
29+
result = self.api.search_indicators(name="test")
30+
self.assertEqual(result, [{"name": "test"}])
31+
mock_post.assert_called_with(
32+
"http://fake-url/api/v2/indicators/search",
33+
json={"query": {"name": "test"}, "count": 0},
34+
)
35+
36+
@patch("yeti.api.requests.Session.post")
37+
def test_search_entities(self, mock_post):
38+
mock_response = MagicMock()
39+
mock_response.json.return_value = {"entities": [{"name": "test_entity"}]}
40+
mock_post.return_value = mock_response
41+
42+
result = self.api.search_entities(name="test_entity")
43+
self.assertEqual(result, [{"name": "test_entity"}])
44+
mock_post.assert_called_with(
45+
"http://fake-url/api/v2/entities/search",
46+
json={"query": {"name": "test_entity"}, "count": 0},
47+
)
48+
49+
@patch("yeti.api.requests.Session.post")
50+
def test_search_observables(self, mock_post):
51+
mock_response = MagicMock()
52+
mock_response.json.return_value = {"observables": [{"value": "test_value"}]}
53+
mock_post.return_value = mock_response
54+
55+
result = self.api.search_observables(value="test_value")
56+
self.assertEqual(result, [{"value": "test_value"}])
57+
mock_post.assert_called_with(
58+
"http://fake-url/api/v2/observables/search",
59+
json={"query": {"value": "test_value"}, "count": 0},
60+
)
61+
62+
@patch("yeti.api.requests.Session.post")
63+
def test_new_entity(self, mock_post):
64+
mock_response = MagicMock()
65+
mock_response.json.return_value = {"id": "new_entity"}
66+
mock_post.return_value = mock_response
67+
68+
result = self.api.new_entity({"name": "test_entity"})
69+
self.assertEqual(result, {"id": "new_entity"})
70+
mock_post.assert_called_with(
71+
"http://fake-url/api/v2/entities/",
72+
json={"entity": {"name": "test_entity"}},
73+
)
74+
75+
@patch("yeti.api.requests.Session.post")
76+
def test_new_indicator(self, mock_post):
77+
mock_response = MagicMock()
78+
mock_response.json.return_value = {"id": "new_indicator"}
79+
mock_post.return_value = mock_response
80+
81+
result = self.api.new_indicator({"name": "test_indicator"})
82+
self.assertEqual(result, {"id": "new_indicator"})
83+
mock_post.assert_called_with(
84+
"http://fake-url/api/v2/indicators/",
85+
json={"indicator": {"name": "test_indicator"}},
86+
)
87+
88+
@patch("yeti.api.requests.Session.patch")
89+
def test_patch_indicator(self, mock_patch):
90+
mock_response = MagicMock()
91+
mock_response.json.return_value = {"id": "patched_indicator"}
92+
mock_patch.return_value = mock_response
93+
94+
result = self.api.patch_indicator(1, {"name": "patched_indicator"})
95+
self.assertEqual(result, {"id": "patched_indicator"})
96+
mock_patch.assert_called_with(
97+
"http://fake-url/api/v2/indicators/1",
98+
json={"indicator": {"name": "patched_indicator"}},
99+
)
100+
101+
@patch("yeti.api.requests.Session.post")
102+
def test_search_dfiq(self, mock_post):
103+
mock_response = MagicMock()
104+
mock_response.json.return_value = {"dfiq": [{"name": "test_dfiq"}]}
105+
mock_post.return_value = mock_response
106+
107+
result = self.api.search_dfiq(name="test_dfiq")
108+
self.assertEqual(result, [{"name": "test_dfiq"}])
109+
mock_post.assert_called_with(
110+
"http://fake-url/api/v2/dfiq/search",
111+
json={"query": {"name": "test_dfiq"}, "count": 0},
112+
)
113+
114+
@patch("yeti.api.requests.Session.post")
115+
def test_new_dfiq_from_yaml(self, mock_post):
116+
mock_response = MagicMock()
117+
mock_response.json.return_value = {"id": "new_dfiq"}
118+
mock_post.return_value = mock_response
119+
120+
result = self.api.new_dfiq_from_yaml("type", "yaml_content")
121+
self.assertEqual(result, {"id": "new_dfiq"})
122+
mock_post.assert_called_with(
123+
"http://fake-url/api/v2/dfiq/from_yaml",
124+
json={
125+
"dfiq_type": "type",
126+
"dfiq_yaml": "yaml_content",
127+
"update_indicators": True,
128+
},
129+
)
130+
131+
@patch("yeti.api.requests.Session.patch")
132+
def test_patch_dfiq_from_yaml(self, mock_patch):
133+
mock_response = MagicMock()
134+
mock_response.json.return_value = {"id": "patched_dfiq"}
135+
mock_patch.return_value = mock_response
136+
137+
result = self.api.patch_dfiq_from_yaml("type", "yaml_content", 1)
138+
self.assertEqual(result, {"id": "patched_dfiq"})
139+
mock_patch.assert_called_with(
140+
"http://fake-url/api/v2/dfiq/1",
141+
json={
142+
"dfiq_type": "type",
143+
"dfiq_yaml": "yaml_content",
144+
"update_indicators": True,
145+
},
146+
)
147+
148+
@patch("yeti.api.requests.Session.post")
149+
def test_download_dfiq_archive(self, mock_post):
150+
mock_response = MagicMock()
151+
mock_response.bytes = b"archive_content"
152+
mock_post.return_value = mock_response
153+
154+
result = self.api.download_dfiq_archive()
155+
self.assertEqual(result, b"archive_content")
156+
mock_post.assert_called_with(
157+
"http://fake-url/api/v2/dfiq/to_archive",
158+
json={"count": 0},
159+
)
160+
161+
@patch("yeti.api.requests.Session.post")
162+
def test_upload_dfiq_archive(self, mock_post):
163+
mock_response = MagicMock()
164+
mock_response.json.return_value = {"uploaded": 1}
165+
mock_post.return_value = mock_response
166+
167+
with patch("builtins.open", unittest.mock.mock_open(read_data=b"data")):
168+
result = self.api.upload_dfiq_archive("path/to/archive.zip")
169+
self.assertEqual(result, {"uploaded": 1})
170+
self.assertEqual(
171+
mock_post.call_args[0][0], "http://fake-url/api/v2/dfiq/from_archive"
172+
)
173+
self.assertRegex(
174+
mock_post.call_args[1]["extra_headers"]["Content-Type"],
175+
"multipart/form-data; boundary=[a-f0-9]{32}",
176+
)
177+
178+
@patch("yeti.api.requests.Session.post")
179+
def test_add_observable(self, mock_post):
180+
mock_response = MagicMock()
181+
mock_response.json.return_value = {"id": "new_observable"}
182+
mock_post.return_value = mock_response
183+
184+
result = self.api.add_observable("value", "type")
185+
self.assertEqual(result, {"id": "new_observable"})
186+
mock_post.assert_called_with(
187+
"http://fake-url/api/v2/observables/",
188+
json={"value": "value", "type": "type", "tags": None},
189+
)
190+
191+
@patch("yeti.api.requests.Session.post")
192+
def test_add_observables_bulk(self, mock_post):
193+
mock_response = MagicMock()
194+
mock_response.json.return_value = {"added": [], "failed": []}
195+
mock_post.return_value = mock_response
196+
197+
result = self.api.add_observables_bulk([{"value": "value", "type": "type"}])
198+
self.assertEqual(result, {"added": [], "failed": []})
199+
mock_post.assert_called_with(
200+
"http://fake-url/api/v2/observables/bulk",
201+
json={"observables": [{"value": "value", "type": "type"}]},
202+
)
203+
204+
@patch("yeti.api.requests.Session.post")
205+
def test_tag_object(self, mock_post):
206+
mock_response = MagicMock()
207+
mock_response.json.return_value = {"id": "tagged_object"}
208+
mock_post.return_value = mock_response
209+
210+
result = self.api.tag_object({"id": "1", "root_type": "indicator"}, ["tag1"])
211+
self.assertEqual(result, {"id": "tagged_object"})
212+
mock_post.assert_called_with(
213+
"http://fake-url/api/v2/indicators/tag",
214+
json={"tags": ["tag1"], "ids": ["1"]},
215+
)
216+
217+
@patch("yeti.api.requests.Session.post")
218+
def test_link_objects(self, mock_post):
219+
mock_response = MagicMock()
220+
mock_response.json.return_value = {"id": "link"}
221+
mock_post.return_value = mock_response
222+
223+
result = self.api.link_objects(
224+
{"id": "1", "root_type": "indicator"},
225+
{"id": "2", "root_type": "entity"},
226+
"link_type",
227+
)
228+
self.assertEqual(result, {"id": "link"})
229+
mock_post.assert_called_with(
230+
"http://fake-url/api/v2/graph/add",
231+
json={
232+
"source": "indicator/1",
233+
"target": "entity/2",
234+
"link_type": "link_type",
235+
"description": None,
236+
},
237+
)
238+
239+
@patch("yeti.api.requests.Session.post")
240+
def test_search_graph(self, mock_post):
241+
mock_response = MagicMock()
242+
mock_response.json.return_value = {"graph": "data"}
243+
mock_post.return_value = mock_response
244+
245+
result = self.api.search_graph("source", "graph", ["type"])
246+
self.assertEqual(result, {"graph": "data"})
247+
mock_post.assert_called_with(
248+
"http://fake-url/api/v2/graph/search",
249+
json={
250+
"count": 0,
251+
"source": "source",
252+
"graph": "graph",
253+
"min_hops": 1,
254+
"max_hops": 1,
255+
"direction": "outbound",
256+
"include_original": True,
257+
"target_types": ["type"],
258+
},
259+
)
260+
261+
262+
if __name__ == "__main__":
263+
unittest.main()

‎yeti/api.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Python client for the Yeti API."""
22

33
import requests
4+
import requests_toolbelt.multipart.encoder as encoder
45

56
import json
67
from typing import Any, Sequence
@@ -149,10 +150,8 @@ def new_indicator(
149150
The response from the API; a dict representing the indicator.
150151
"""
151152
params = {"indicator": indicator}
152-
response, _ = self.client.post(
153-
f"{self._url_root}/api/v2/indicators/", json=params
154-
)
155-
indicator = json.loads(response)
153+
response = self.client.post(f"{self._url_root}/api/v2/indicators/", json=params)
154+
indicator = response.json()
156155

157156
if tags:
158157
params = {"tags": tags, "ids": [indicator["id"]]}
@@ -235,7 +234,7 @@ def download_dfiq_archive(self, dfiq_type: str | None = None) -> bytes:
235234
response = self.client.post(
236235
f"{self._url_root}/api/v2/dfiq/to_archive", json=params
237236
)
238-
return body
237+
return response.bytes
239238

240239
def upload_dfiq_archive(self, archive_path: str) -> dict[str, int]:
241240
"""Uploads a DFIQ archive to Yeti.
@@ -318,8 +317,8 @@ def tag_object(
318317
"""Tags an object in Yeti."""
319318
params = {"tags": list(tags), "ids": [yeti_object["id"]]}
320319
endpoint = TYPE_TO_ENDPOINT[yeti_object["root_type"]]
321-
result, _ = self.client.post(f"{self._url_root}{endpoint}/tag", json=params)
322-
return json.loads(result)
320+
result = self.client.post(f"{self._url_root}{endpoint}/tag", json=params)
321+
return result.json()
323322

324323
def link_objects(
325324
self,

0 commit comments

Comments
 (0)