Support max num results

Author: adamtheturtle

src/vws/query.py

@@ -32,12 +32,16 @@ def __init__(
     def query(
         self,
         image: io.BytesIO,
+        max_num_results: int = 1,
     ) -> str:
         """
         TODO docstring
         """
         image_content = image.getvalue()
-        body = {'image': ('image.jpeg', image_content, 'image/jpeg')}
+        body = {
+            'image': ('image.jpeg', image_content, 'image/jpeg'),
+            'max_num_results': (None, max_num_results, 'text/plain'),
+        }
         date = rfc_1123_date()
         request_path = '/v1/query'
         content, content_type_header = encode_multipart_formdata(body)

tests/test_query.py

@@ -3,6 +3,7 @@
 """
 
 import io
+import uuid
 
 from vws import VWS, CloudRecoService
 
@@ -42,6 +43,79 @@ def test_match(
         assert matching_target['target_id'] == target_id
 
 
+class TestMaxNumResults:
+    """
+    Tests for the ``max_num_results`` parameter of ``query``.
+    """
+
+    def test_default(
+        self,
+        client: VWS,
+        cloud_reco_client: CloudRecoService,
+        high_quality_image: io.BytesIO,
+    ) -> None:
+        """
+        XXX
+        """
+        target_id = client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        target_id_2 = client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        client.wait_for_target_processed(target_id=target_id)
+        client.wait_for_target_processed(target_id=target_id_2)
+        matches = cloud_reco_client.query(image=high_quality_image)
+        assert len(matches) == 1
+
+    def test_custom(
+        self,
+        client: VWS,
+        cloud_reco_client: CloudRecoService,
+        high_quality_image: io.BytesIO,
+    ) -> None:
+        """
+        XXX
+        """
+        target_id = client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        target_id_2 = client.add_target(
+            name=uuid.uuid4().hex,
+            width=1,
+            image=high_quality_image,
+        )
+        client.wait_for_target_processed(target_id=target_id)
+        client.wait_for_target_processed(target_id=target_id_2)
+        matches = cloud_reco_client.query(image=high_quality_image, max_num_results=2)
+        assert len(matches) == 2
+
+
+
+    def test_foo(self):
+        pass
+        # target_ids = set([])
+        # for i in range(15):
+        #     target_id = client.add_target(
+        #         name=uuid.uuid4().hex,
+        #         width=1,
+        #         image=high_quality_image,
+        #     )
+        #     target_ids.add(target_id)
+        #
+        # for target_id in target_ids:
+        #     client.wait_for_target_processed(target_id=target_id)
+        #
+        # matching_targets = cloud_reco_client.query(image=high_quality_image)
+        # assert len(matching_targets) == 1
+
+
 # TODO test custom base URL
 # TODO test bad credentials
 # TODO test options - max_num_results + include_target_data