Tweak some tests

log2timeline · Dec 8, 2023 · dfb17f1 · dfb17f1
1 parent 9a8c7a6
commit dfb17f1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/dftimewolf/lib/processors/turbinia_base.py b/dftimewolf/lib/processors/turbinia_base.py
@@ -275,7 +275,9 @@ def RefreshClientCredentials(self) -> bool:
     return bool(refresh)
 
   def InitializeTurbiniaApiClient(
-      self, credentials: Credentials) -> turbinia_api_lib.api_client.ApiClient:
+      self,
+      credentials: Optional[Credentials]
+    ) -> turbinia_api_lib.api_client.ApiClient:
     """Creates a Turbinia API client object.
 
     This method also runs the authentication flow if needed.

diff --git a/tests/lib/processors/turbinia_base.py b/tests/lib/processors/turbinia_base.py
@@ -58,14 +58,14 @@ def testTurbiniaSetup(self, _mock_read_config):
         turbinia_zone="us-central1f",
         turbinia_api="http://localhost:8001",
         incident_id="123456789",
-        sketch_id="12345",
+        sketch_id=12345,
     )
     self.assertEqual(self.turbinia_processor.project, "turbinia-project")
     self.assertEqual(self.turbinia_processor.turbinia_zone, "us-central1f")
     self.assertEqual(
         self.turbinia_processor.turbinia_api, "http://localhost:8001")
     self.assertEqual(self.turbinia_processor.incident_id, "123456789")
-    self.assertEqual(self.turbinia_processor.sketch_id, "12345")
+    self.assertEqual(self.turbinia_processor.sketch_id, 12345)
     self.assertEqual(self.turbinia_processor.output_path, "/tmp")
     self.assertEqual(self.turbinia_processor.turbinia_recipe, None)
 
@@ -88,7 +88,7 @@ def testTurbiniaStart(self, mock_create_request):
         "zone": "us-central1-f",
     }
     request_id = self.turbinia_processor.TurbiniaStart(
-        evidence=evidence, yara_rules=YARA_RULE)
+        evidence=evidence, threat_intel_indicators=None, yara_rules=YARA_RULE)
     self.assertEqual(request_id, "41483253079448e59685d88f37ab91f7")
 
   # pylint: disable=line-too-long