from __future__ import print_function
import json
import unittest

from censys import CensysAPIBase, CensysException


class CensysIPv4(CensysAPIBase):

    def search(self, query, fields=None):
        if fields is None:
            fields = []
        page = 1
        pages = float('inf')
        data = {
            "query": query,
            "page": page,
            "fields": fields
        }

        while page <= pages:
            payload = self._post("search/ipv4", data=data)
            pages = payload['metadata']['pages']
            page += 1

            for result in payload["results"]:
                yield result

    def view(self, ip):
        return self._get("/".join(("view", "ipv4", ip)))

    def report(self, query, field, buckets=50):
        data = {
            "query": query,
            "field": field,
            "buckets": int(buckets)
        }
        return self._post("report/ipv4", data=data)


class CensysIPv4Tests(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls._api = CensysIPv4()

    def testGet(self):
        print(self._api.view("84.206.102.184"))

    def testEmptySearch(self):
        with self.assertRaises(CensysException):
            self._api._post("search/ipv4", data={"query1": "query"})

    def testSearch(self):
        print(json.dumps(self._api.search("*")))

    def testReport(self):
        print(self._api.report("*", "protocols", 5))


if __name__ == "__main__":
    unittest.main()

@@ -1313,25 +1313,16 @@
     
-print(json.dumps(
+hosts = 
 self
@@ -1338,16 +1338,68 @@
 rch(%22*%22)
+%0A        print(next(hosts))%0A        print(next(hosts
 ))%0A%0A    

fix test for generator